initial import

core/discov/internal/registry_test.go

@@ -0,0 +1,245 @@
+package internal
+
+import (
+	"context"
+	"sync"
+	"testing"
+
+	"zero/core/contextx"
+	"zero/core/stringx"
+
+	"github.com/coreos/etcd/mvcc/mvccpb"
+	"github.com/golang/mock/gomock"
+	"github.com/stretchr/testify/assert"
+	"go.etcd.io/etcd/clientv3"
+)
+
+var mockLock sync.Mutex
+
+func setMockClient(cli EtcdClient) func() {
+	mockLock.Lock()
+	NewClient = func([]string) (EtcdClient, error) {
+		return cli, nil
+	}
+	return func() {
+		NewClient = DialClient
+		mockLock.Unlock()
+	}
+}
+
+func TestGetCluster(t *testing.T) {
+	c1 := GetRegistry().getCluster([]string{"first"})
+	c2 := GetRegistry().getCluster([]string{"second"})
+	c3 := GetRegistry().getCluster([]string{"first"})
+	assert.Equal(t, c1, c3)
+	assert.NotEqual(t, c1, c2)
+}
+
+func TestGetClusterKey(t *testing.T) {
+	assert.Equal(t, getClusterKey([]string{"localhost:1234", "remotehost:5678"}),
+		getClusterKey([]string{"remotehost:5678", "localhost:1234"}))
+}
+
+func TestCluster_HandleChanges(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	l := NewMockUpdateListener(ctrl)
+	l.EXPECT().OnAdd(KV{
+		Key: "first",
+		Val: "1",
+	})
+	l.EXPECT().OnAdd(KV{
+		Key: "second",
+		Val: "2",
+	})
+	l.EXPECT().OnDelete(KV{
+		Key: "first",
+		Val: "1",
+	})
+	l.EXPECT().OnDelete(KV{
+		Key: "second",
+		Val: "2",
+	})
+	l.EXPECT().OnAdd(KV{
+		Key: "third",
+		Val: "3",
+	})
+	l.EXPECT().OnAdd(KV{
+		Key: "fourth",
+		Val: "4",
+	})
+	c := newCluster([]string{"any"})
+	c.listeners["any"] = []UpdateListener{l}
+	c.handleChanges("any", []KV{
+		{
+			Key: "first",
+			Val: "1",
+		},
+		{
+			Key: "second",
+			Val: "2",
+		},
+	})
+	assert.EqualValues(t, map[string]string{
+		"first":  "1",
+		"second": "2",
+	}, c.values["any"])
+	c.handleChanges("any", []KV{
+		{
+			Key: "third",
+			Val: "3",
+		},
+		{
+			Key: "fourth",
+			Val: "4",
+		},
+	})
+	assert.EqualValues(t, map[string]string{
+		"third":  "3",
+		"fourth": "4",
+	}, c.values["any"])
+}
+
+func TestCluster_Load(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	cli := NewMockEtcdClient(ctrl)
+	restore := setMockClient(cli)
+	defer restore()
+	cli.EXPECT().Get(gomock.Any(), "any/", gomock.Any()).Return(&clientv3.GetResponse{
+		Kvs: []*mvccpb.KeyValue{
+			{
+				Key:   []byte("hello"),
+				Value: []byte("world"),
+			},
+		},
+	}, nil)
+	cli.EXPECT().Ctx().Return(context.Background())
+	c := &cluster{
+		values: make(map[string]map[string]string),
+	}
+	c.load(cli, "any")
+}
+
+func TestCluster_Watch(t *testing.T) {
+	tests := []struct {
+		name      string
+		method    int
+		eventType mvccpb.Event_EventType
+	}{
+		{
+			name:      "add",
+			eventType: clientv3.EventTypePut,
+		},
+		{
+			name:      "delete",
+			eventType: clientv3.EventTypeDelete,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			ctrl := gomock.NewController(t)
+			defer ctrl.Finish()
+			cli := NewMockEtcdClient(ctrl)
+			restore := setMockClient(cli)
+			defer restore()
+			ch := make(chan clientv3.WatchResponse)
+			cli.EXPECT().Watch(gomock.Any(), "any/", gomock.Any()).Return(ch)
+			cli.EXPECT().Ctx().Return(context.Background())
+			var wg sync.WaitGroup
+			wg.Add(1)
+			c := &cluster{
+				listeners: make(map[string][]UpdateListener),
+				values:    make(map[string]map[string]string),
+			}
+			listener := NewMockUpdateListener(ctrl)
+			c.listeners["any"] = []UpdateListener{listener}
+			listener.EXPECT().OnAdd(gomock.Any()).Do(func(kv KV) {
+				assert.Equal(t, "hello", kv.Key)
+				assert.Equal(t, "world", kv.Val)
+				wg.Done()
+			}).MaxTimes(1)
+			listener.EXPECT().OnDelete(gomock.Any()).Do(func(_ interface{}) {
+				wg.Done()
+			}).MaxTimes(1)
+			go c.watch(cli, "any")
+			ch <- clientv3.WatchResponse{
+				Events: []*clientv3.Event{
+					{
+						Type: test.eventType,
+						Kv: &mvccpb.KeyValue{
+							Key:   []byte("hello"),
+							Value: []byte("world"),
+						},
+					},
+				},
+			}
+			wg.Wait()
+		})
+	}
+}
+
+func TestClusterWatch_RespFailures(t *testing.T) {
+	resps := []clientv3.WatchResponse{
+		{
+			Canceled: true,
+		},
+		{
+			// cause resp.Err() != nil
+			CompactRevision: 1,
+		},
+	}
+
+	for _, resp := range resps {
+		t.Run(stringx.Rand(), func(t *testing.T) {
+			ctrl := gomock.NewController(t)
+			defer ctrl.Finish()
+			cli := NewMockEtcdClient(ctrl)
+			restore := setMockClient(cli)
+			defer restore()
+			ch := make(chan clientv3.WatchResponse)
+			cli.EXPECT().Watch(gomock.Any(), "any/", gomock.Any()).Return(ch)
+			cli.EXPECT().Ctx().Return(context.Background()).AnyTimes()
+			c := new(cluster)
+			go func() {
+				ch <- resp
+			}()
+			c.watch(cli, "any")
+		})
+	}
+}
+
+func TestClusterWatch_CloseChan(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	cli := NewMockEtcdClient(ctrl)
+	restore := setMockClient(cli)
+	defer restore()
+	ch := make(chan clientv3.WatchResponse)
+	cli.EXPECT().Watch(gomock.Any(), "any/", gomock.Any()).Return(ch)
+	cli.EXPECT().Ctx().Return(context.Background()).AnyTimes()
+	c := new(cluster)
+	go func() {
+		close(ch)
+	}()
+	c.watch(cli, "any")
+}
+
+func TestValueOnlyContext(t *testing.T) {
+	ctx := contextx.ValueOnlyFrom(context.Background())
+	ctx.Done()
+	assert.Nil(t, ctx.Err())
+}
+
+type mockedSharedCalls struct {
+	fn func() (interface{}, error)
+}
+
+func (c mockedSharedCalls) Do(_ string, fn func() (interface{}, error)) (interface{}, error) {
+	return c.fn()
+}
+
+func (c mockedSharedCalls) DoEx(_ string, fn func() (interface{}, error)) (interface{}, bool, error) {
+	val, err := c.fn()
+	return val, true, err
+}