kratos/pkg/net/rpc/warden/resolver/test/mockdiscovery.go

package resolver

import (
	"context"

	"github.com/go-kratos/kratos/pkg/conf/env"
	"github.com/go-kratos/kratos/pkg/naming"
)

type mockDiscoveryBuilder struct {
	instances map[string]*naming.Instance
	watchch   map[string][]*mockDiscoveryResolver
}

func (mb *mockDiscoveryBuilder) Build(id string, opts ...naming.BuildOpt) 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
		}
	}
}