package resolver import ( "context" "github.com/bilibili/kratos/pkg/conf/env" "github.com/bilibili/kratos/pkg/naming" ) type mockDiscoveryBuilder struct { instances map[string]*naming.Instance watchch map[string][]*mockDiscoveryResolver } func (mb *mockDiscoveryBuilder) Build(id string) naming.Resolver { mr := &mockDiscoveryResolver{ d: mb, watchch: make(chan struct{}, 1), } mb.watchch[id] = append(mb.watchch[id], mr) mr.watchch <- struct{}{} return mr } func (mb *mockDiscoveryBuilder) Scheme() string { return "mockdiscovery" } type mockDiscoveryResolver struct { d *mockDiscoveryBuilder watchch chan struct{} } var _ naming.Resolver = &mockDiscoveryResolver{} func (md *mockDiscoveryResolver) Fetch(ctx context.Context) (*naming.InstancesInfo, bool) { zones := make(map[string][]*naming.Instance) for _, v := range md.d.instances { zones[v.Zone] = append(zones[v.Zone], v) } return &naming.InstancesInfo{Instances: zones}, len(zones) > 0 } func (md *mockDiscoveryResolver) Watch() <-chan struct{} { return md.watchch } func (md *mockDiscoveryResolver) Close() error { close(md.watchch) return nil } func (md *mockDiscoveryResolver) Scheme() string { return "mockdiscovery" } func (mb *mockDiscoveryBuilder) registry(appID string, hostname, rpc string, metadata map[string]string) { ins := &naming.Instance{ AppID: appID, Env: "hello=world", Hostname: hostname, Addrs: []string{"grpc://" + rpc}, Version: "1.1", Zone: env.Zone, Metadata: metadata, } mb.instances[hostname] = ins if ch, ok := mb.watchch[appID]; ok { var bullet struct{} for _, c := range ch { c.watchch <- bullet } } } func (mb *mockDiscoveryBuilder) cancel(hostname string) { ins, ok := mb.instances[hostname] if !ok { return } delete(mb.instances, hostname) if ch, ok := mb.watchch[ins.AppID]; ok { var bullet struct{} for _, c := range ch { c.watchch <- bullet } } }