add ut

Signed-off-by: Wei Liu <[email protected]>

internal/proxy/shard_client_test.go

@@ -1,19 +1,66 @@
 package proxy
 
-import "testing"
-
-func genShardLeaderInfo(channel string, leaderIDs []UniqueID) map[string][]nodeInfo {
-	leaders := make(map[string][]nodeInfo)
-	nodeInfos := make([]nodeInfo, len(leaderIDs))
-	for i, id := range leaderIDs {
-		nodeInfos[i] = nodeInfo{
-			nodeID:  id,
-			address: "fake",
-		}
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/milvus-io/milvus/internal/mocks"
+	"github.com/milvus-io/milvus/internal/types"
+	"github.com/milvus-io/milvus/pkg/util/paramtable"
+)
+
+func TestShardClientMgr(t *testing.T) {
+	ctx := context.Background()
+	nodeInfo := nodeInfo{
+		nodeID: 1,
 	}
-	leaders[channel] = nodeInfos
-	return leaders
+
+	qn := mocks.NewMockQueryNodeClient(t)
+	qn.EXPECT().Close().Return(nil)
+	creator := func(ctx context.Context, addr string, nodeID int64) (types.QueryNodeClient, error) {
+		return qn, nil
+	}
+
+	mgr := newShardClientMgr()
+	mgr.SetClientCreatorFunc(creator)
+	_, err := mgr.GetClient(ctx, nodeInfo)
+	assert.Nil(t, err)
+
+	mgr.ReleaseClient(1)
+	assert.Equal(t, len(mgr.clients.data), 1)
+	mgr.Close()
+	assert.Equal(t, len(mgr.clients.data), 0)
 }
 
-func TestShardClientMgr(t *testing.T) {
+func TestShardClient(t *testing.T) {
+	nodeInfo := nodeInfo{
+		nodeID: 1,
+	}
+
+	qn := mocks.NewMockQueryNodeClient(t)
+	qn.EXPECT().Close().Return(nil)
+	creator := func(ctx context.Context, addr string, nodeID int64) (types.QueryNodeClient, error) {
+		return qn, nil
+	}
+	shardClient, err := newShardClient(nodeInfo, creator)
+	assert.Nil(t, err)
+	assert.Equal(t, len(shardClient.clients), 0)
+	assert.Equal(t, int64(1), shardClient.refCnt.Load())
+	assert.Equal(t, false, shardClient.initialized.Load())
+
+	ctx := context.Background()
+	_, err = shardClient.getClient(ctx)
+	assert.Nil(t, err)
+	assert.Equal(t, len(shardClient.clients), paramtable.Get().ProxyCfg.QueryNodePoolingSize.GetAsInt())
+	assert.Equal(t, int64(2), shardClient.refCnt.Load())
+	assert.Equal(t, true, shardClient.initialized.Load())
+
+	shardClient.Release()
+	assert.Equal(t, int64(1), shardClient.refCnt.Load())
+
+	shardClient.Release()
+	assert.Equal(t, int64(0), shardClient.refCnt.Load())
+	assert.Equal(t, true, shardClient.isClosed)
 }