From db3c6814f4bdaf86f979c8cc089486950bd7126c Mon Sep 17 00:00:00 2001 From: Congqi Xia Date: Mon, 12 Aug 2024 18:05:11 +0800 Subject: [PATCH] Fix proxy unit tests Signed-off-by: Congqi Xia --- internal/proxy/impl_test.go | 2 ++ internal/proxy/lb_policy_test.go | 1 + internal/proxy/meta_cache_test.go | 12 +++++++ internal/proxy/msg_pack_test.go | 6 +++- internal/proxy/task_index_test.go | 45 +++++++++++++++++++++++++++ internal/proxy/task_query_test.go | 3 ++ internal/proxy/task_search_test.go | 6 ++++ internal/proxy/task_statistic_test.go | 1 + internal/proxy/task_test.go | 21 ++++++++++++- 9 files changed, 95 insertions(+), 2 deletions(-) diff --git a/internal/proxy/impl_test.go b/internal/proxy/impl_test.go index c9f7f2439ef3f..6104b99be5688 100644 --- a/internal/proxy/impl_test.go +++ b/internal/proxy/impl_test.go @@ -249,6 +249,8 @@ func TestProxy_ResourceGroup(t *testing.T) { qc := mocks.NewMockQueryCoordClient(t) node.SetQueryCoordClient(qc) + qc.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{}, nil).Maybe() + tsoAllocatorIns := newMockTsoAllocator() node.sched, err = newTaskScheduler(node.ctx, tsoAllocatorIns, node.factory) assert.NoError(t, err) diff --git a/internal/proxy/lb_policy_test.go b/internal/proxy/lb_policy_test.go index 86551df24a6d5..a83e1a5c068b0 100644 --- a/internal/proxy/lb_policy_test.go +++ b/internal/proxy/lb_policy_test.go @@ -67,6 +67,7 @@ func (s *LBPolicySuite) SetupTest() { successStatus := commonpb.Status{ErrorCode: commonpb.ErrorCode_Success} qc := mocks.NewMockQueryCoordClient(s.T()) qc.EXPECT().LoadCollection(mock.Anything, mock.Anything).Return(&successStatus, nil) + qc.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{}, nil).Maybe() qc.EXPECT().GetShardLeaders(mock.Anything, mock.Anything).Return(&querypb.GetShardLeadersResponse{ Status: &successStatus, diff --git a/internal/proxy/meta_cache_test.go b/internal/proxy/meta_cache_test.go index f2459b674af70..8a98dd8d6476d 100644 --- a/internal/proxy/meta_cache_test.go +++ b/internal/proxy/meta_cache_test.go @@ -195,6 +195,9 @@ func TestMetaCache_GetCollection(t *testing.T) { ctx := context.Background() rootCoord := &MockRootCoordClientInterface{} queryCoord := &mocks.MockQueryCoordClient{} + + queryCoord.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{}, nil).Maybe() + mgr := newShardClientMgr() err := InitMetaCache(ctx, rootCoord, queryCoord, mgr) assert.NoError(t, err) @@ -245,6 +248,8 @@ func TestMetaCache_GetBasicCollectionInfo(t *testing.T) { ctx := context.Background() rootCoord := &MockRootCoordClientInterface{} queryCoord := &mocks.MockQueryCoordClient{} + + queryCoord.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{}, nil).Maybe() mgr := newShardClientMgr() err := InitMetaCache(ctx, rootCoord, queryCoord, mgr) assert.NoError(t, err) @@ -277,6 +282,7 @@ func TestMetaCache_GetCollectionName(t *testing.T) { ctx := context.Background() rootCoord := &MockRootCoordClientInterface{} queryCoord := &mocks.MockQueryCoordClient{} + queryCoord.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{}, nil).Maybe() mgr := newShardClientMgr() err := InitMetaCache(ctx, rootCoord, queryCoord, mgr) assert.NoError(t, err) @@ -327,6 +333,7 @@ func TestMetaCache_GetCollectionFailure(t *testing.T) { ctx := context.Background() rootCoord := &MockRootCoordClientInterface{} queryCoord := &mocks.MockQueryCoordClient{} + queryCoord.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{}, nil).Maybe() mgr := newShardClientMgr() err := InitMetaCache(ctx, rootCoord, queryCoord, mgr) assert.NoError(t, err) @@ -360,6 +367,7 @@ func TestMetaCache_GetNonExistCollection(t *testing.T) { ctx := context.Background() rootCoord := &MockRootCoordClientInterface{} queryCoord := &mocks.MockQueryCoordClient{} + queryCoord.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{}, nil).Maybe() mgr := newShardClientMgr() err := InitMetaCache(ctx, rootCoord, queryCoord, mgr) assert.NoError(t, err) @@ -376,6 +384,7 @@ func TestMetaCache_GetPartitionID(t *testing.T) { ctx := context.Background() rootCoord := &MockRootCoordClientInterface{} queryCoord := &mocks.MockQueryCoordClient{} + queryCoord.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{}, nil).Maybe() mgr := newShardClientMgr() err := InitMetaCache(ctx, rootCoord, queryCoord, mgr) assert.NoError(t, err) @@ -398,6 +407,7 @@ func TestMetaCache_ConcurrentTest1(t *testing.T) { ctx := context.Background() rootCoord := &MockRootCoordClientInterface{} queryCoord := &mocks.MockQueryCoordClient{} + queryCoord.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{}, nil).Maybe() mgr := newShardClientMgr() err := InitMetaCache(ctx, rootCoord, queryCoord, mgr) assert.NoError(t, err) @@ -452,6 +462,7 @@ func TestMetaCache_GetPartitionError(t *testing.T) { ctx := context.Background() rootCoord := &MockRootCoordClientInterface{} queryCoord := &mocks.MockQueryCoordClient{} + queryCoord.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{}, nil).Maybe() mgr := newShardClientMgr() err := InitMetaCache(ctx, rootCoord, queryCoord, mgr) assert.NoError(t, err) @@ -805,6 +816,7 @@ func TestMetaCache_Database(t *testing.T) { ctx := context.Background() rootCoord := &MockRootCoordClientInterface{} queryCoord := &mocks.MockQueryCoordClient{} + queryCoord.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{}, nil).Maybe() shardMgr := newShardClientMgr() err := InitMetaCache(ctx, rootCoord, queryCoord, shardMgr) assert.NoError(t, err) diff --git a/internal/proxy/msg_pack_test.go b/internal/proxy/msg_pack_test.go index 29873f11b0821..9d199dd166171 100644 --- a/internal/proxy/msg_pack_test.go +++ b/internal/proxy/msg_pack_test.go @@ -29,6 +29,8 @@ import ( "github.com/milvus-io/milvus-proto/go-api/v2/msgpb" "github.com/milvus-io/milvus-proto/go-api/v2/schemapb" "github.com/milvus-io/milvus/internal/allocator" + "github.com/milvus-io/milvus/internal/mocks" + "github.com/milvus-io/milvus/internal/proto/querypb" "github.com/milvus-io/milvus/pkg/mq/msgstream" "github.com/milvus-io/milvus/pkg/util/funcutil" "github.com/milvus-io/milvus/pkg/util/paramtable" @@ -153,8 +155,10 @@ func TestRepackInsertDataWithPartitionKey(t *testing.T) { rc := NewRootCoordMock() defer rc.Close() + qc := &mocks.MockQueryCoordClient{} + qc.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{}, nil).Maybe() - err := InitMetaCache(ctx, rc, nil, nil) + err := InitMetaCache(ctx, rc, qc, nil) assert.NoError(t, err) idAllocator, err := allocator.NewIDAllocator(ctx, rc, paramtable.GetNodeID()) diff --git a/internal/proxy/task_index_test.go b/internal/proxy/task_index_test.go index 47251cc2edf72..4f709f2a9c77d 100644 --- a/internal/proxy/task_index_test.go +++ b/internal/proxy/task_index_test.go @@ -188,6 +188,18 @@ func TestDropIndexTask_PreExecute(t *testing.T) { t.Run("coll has been loaded", func(t *testing.T) { qc := getMockQueryCoord() + qc.ExpectedCalls = nil + qc.EXPECT().LoadCollection(mock.Anything, mock.Anything).Return(merr.Success(), nil) + qc.EXPECT().GetShardLeaders(mock.Anything, mock.Anything).Return(&querypb.GetShardLeadersResponse{ + Status: merr.Success(), + Shards: []*querypb.ShardLeadersList{ + { + ChannelName: "channel-1", + NodeIds: []int64{1, 2, 3}, + NodeAddrs: []string{"localhost:9000", "localhost:9001", "localhost:9002"}, + }, + }, + }, nil) qc.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{ Status: merr.Success(), CollectionIDs: []int64{collectionID}, @@ -200,6 +212,22 @@ func TestDropIndexTask_PreExecute(t *testing.T) { t.Run("show collection error", func(t *testing.T) { qc := getMockQueryCoord() + qc.ExpectedCalls = nil + qc.EXPECT().LoadCollection(mock.Anything, mock.Anything).Return(merr.Success(), nil) + qc.EXPECT().GetShardLeaders(mock.Anything, mock.Anything).Return(&querypb.GetShardLeadersResponse{ + Status: merr.Success(), + Shards: []*querypb.ShardLeadersList{ + { + ChannelName: "channel-1", + NodeIds: []int64{1, 2, 3}, + NodeAddrs: []string{"localhost:9000", "localhost:9001", "localhost:9002"}, + }, + }, + }, nil) + qc.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{ + Status: merr.Success(), + CollectionIDs: []int64{collectionID}, + }, nil) qc.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(nil, errors.New("error")) dit.queryCoord = qc @@ -209,6 +237,22 @@ func TestDropIndexTask_PreExecute(t *testing.T) { t.Run("show collection fail", func(t *testing.T) { qc := getMockQueryCoord() + qc.ExpectedCalls = nil + qc.EXPECT().LoadCollection(mock.Anything, mock.Anything).Return(merr.Success(), nil) + qc.EXPECT().GetShardLeaders(mock.Anything, mock.Anything).Return(&querypb.GetShardLeadersResponse{ + Status: merr.Success(), + Shards: []*querypb.ShardLeadersList{ + { + ChannelName: "channel-1", + NodeIds: []int64{1, 2, 3}, + NodeAddrs: []string{"localhost:9000", "localhost:9001", "localhost:9002"}, + }, + }, + }, nil) + qc.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{ + Status: merr.Success(), + CollectionIDs: []int64{collectionID}, + }, nil) qc.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{ Status: &commonpb.Status{ ErrorCode: commonpb.ErrorCode_UnexpectedError, @@ -236,6 +280,7 @@ func getMockQueryCoord() *mocks.MockQueryCoordClient { }, }, }, nil) + qc.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{}, nil).Maybe() return qc } diff --git a/internal/proxy/task_query_test.go b/internal/proxy/task_query_test.go index 06d318f6c726a..592a90e0f32b5 100644 --- a/internal/proxy/task_query_test.go +++ b/internal/proxy/task_query_test.go @@ -73,6 +73,9 @@ func TestQueryTask_all(t *testing.T) { }, }, }, nil).Maybe() + qc.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{ + Status: &successStatus, + }, nil).Maybe() mgr := NewMockShardClientManager(t) mgr.EXPECT().GetClient(mock.Anything, mock.Anything).Return(qn, nil).Maybe() diff --git a/internal/proxy/task_search_test.go b/internal/proxy/task_search_test.go index 66cb0dba90cf7..b0f9d769b21aa 100644 --- a/internal/proxy/task_search_test.go +++ b/internal/proxy/task_search_test.go @@ -60,6 +60,9 @@ func TestSearchTask_PostExecute(t *testing.T) { defer rc.Close() require.NoError(t, err) mgr := newShardClientMgr() + qc.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{ + Status: merr.Success(), + }, nil).Maybe() err = InitMetaCache(ctx, rc, qc, mgr) require.NoError(t, err) @@ -191,6 +194,7 @@ func TestSearchTask_PreExecute(t *testing.T) { defer rc.Close() require.NoError(t, err) mgr := newShardClientMgr() + qc.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{}, nil).Maybe() err = InitMetaCache(ctx, rc, qc, mgr) require.NoError(t, err) @@ -335,6 +339,7 @@ func TestSearchTaskV2_Execute(t *testing.T) { defer rc.Close() mgr := newShardClientMgr() + qc.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{}, nil).Maybe() err = InitMetaCache(ctx, rc, qc, mgr) require.NoError(t, err) @@ -1786,6 +1791,7 @@ func TestSearchTask_ErrExecute(t *testing.T) { ) qn.EXPECT().GetComponentStates(mock.Anything, mock.Anything).Return(nil, nil).Maybe() + qc.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{}, nil).Maybe() mgr := NewMockShardClientManager(t) mgr.EXPECT().GetClient(mock.Anything, mock.Anything).Return(qn, nil).Maybe() diff --git a/internal/proxy/task_statistic_test.go b/internal/proxy/task_statistic_test.go index c3438d0f13465..42f0d63b4480d 100644 --- a/internal/proxy/task_statistic_test.go +++ b/internal/proxy/task_statistic_test.go @@ -68,6 +68,7 @@ func (s *StatisticTaskSuite) SetupTest() { }, }, }, nil).Maybe() + qc.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{}, nil).Maybe() qc.EXPECT().ShowPartitions(mock.Anything, mock.Anything).Return(&querypb.ShowPartitionsResponse{ Status: merr.Success(), PartitionIDs: []int64{1, 2, 3}, diff --git a/internal/proxy/task_test.go b/internal/proxy/task_test.go index 06f80220ef588..d70cdc7933736 100644 --- a/internal/proxy/task_test.go +++ b/internal/proxy/task_test.go @@ -977,6 +977,7 @@ func TestHasCollectionTask(t *testing.T) { defer rc.Close() qc := getQueryCoordClient() + qc.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{}, nil).Maybe() ctx := context.Background() mgr := newShardClientMgr() @@ -1123,6 +1124,7 @@ func TestDescribeCollectionTask_ShardsNum1(t *testing.T) { defer rc.Close() qc := getQueryCoordClient() + qc.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{}, nil).Maybe() ctx := context.Background() mgr := newShardClientMgr() @@ -1185,6 +1187,7 @@ func TestDescribeCollectionTask_EnableDynamicSchema(t *testing.T) { rc := NewRootCoordMock() defer rc.Close() qc := getQueryCoordClient() + qc.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{}, nil).Maybe() ctx := context.Background() mgr := newShardClientMgr() InitMetaCache(ctx, rc, qc, mgr) @@ -1248,6 +1251,7 @@ func TestDescribeCollectionTask_ShardsNum2(t *testing.T) { defer rc.Close() qc := getQueryCoordClient() + qc.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{}, nil).Maybe() ctx := context.Background() mgr := newShardClientMgr() @@ -1611,6 +1615,7 @@ func TestTask_Int64PrimaryKey(t *testing.T) { defer rc.Close() qc := getQueryCoordClient() + qc.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{}, nil).Maybe() ctx := context.Background() @@ -1803,6 +1808,7 @@ func TestTask_VarCharPrimaryKey(t *testing.T) { defer rc.Close() qc := getQueryCoordClient() + qc.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{}, nil).Maybe() ctx := context.Background() @@ -2679,6 +2685,7 @@ func TestCreateResourceGroupTask(t *testing.T) { defer rc.Close() qc := getQueryCoordClient() + qc.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{}, nil).Maybe() qc.EXPECT().CreateResourceGroup(mock.Anything, mock.Anything, mock.Anything).Return(merr.Success(), nil) ctx := context.Background() @@ -2719,6 +2726,7 @@ func TestDropResourceGroupTask(t *testing.T) { defer rc.Close() qc := getQueryCoordClient() + qc.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{}, nil).Maybe() qc.EXPECT().DropResourceGroup(mock.Anything, mock.Anything).Return(merr.Success(), nil) ctx := context.Background() @@ -2759,6 +2767,7 @@ func TestTransferNodeTask(t *testing.T) { defer rc.Close() qc := getQueryCoordClient() + qc.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{}, nil).Maybe() qc.EXPECT().TransferNode(mock.Anything, mock.Anything).Return(merr.Success(), nil) ctx := context.Background() @@ -2799,6 +2808,7 @@ func TestTransferNodeTask(t *testing.T) { func TestTransferReplicaTask(t *testing.T) { rc := &MockRootCoordClientInterface{} qc := getQueryCoordClient() + qc.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{}, nil).Maybe() qc.EXPECT().TransferReplica(mock.Anything, mock.Anything).Return(merr.Success(), nil) ctx := context.Background() @@ -2842,6 +2852,7 @@ func TestTransferReplicaTask(t *testing.T) { func TestListResourceGroupsTask(t *testing.T) { rc := &MockRootCoordClientInterface{} qc := getQueryCoordClient() + qc.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{}, nil).Maybe() qc.EXPECT().ListResourceGroups(mock.Anything, mock.Anything).Return(&milvuspb.ListResourceGroupsResponse{ Status: merr.Success(), ResourceGroups: []string{meta.DefaultResourceGroupName, "rg"}, @@ -2885,6 +2896,7 @@ func TestListResourceGroupsTask(t *testing.T) { func TestDescribeResourceGroupTask(t *testing.T) { rc := &MockRootCoordClientInterface{} qc := getQueryCoordClient() + qc.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{}, nil).Maybe() qc.EXPECT().DescribeResourceGroup(mock.Anything, mock.Anything).Return(&querypb.DescribeResourceGroupResponse{ Status: merr.Success(), ResourceGroup: &querypb.ResourceGroupInfo{ @@ -2940,6 +2952,7 @@ func TestDescribeResourceGroupTask(t *testing.T) { func TestDescribeResourceGroupTaskFailed(t *testing.T) { rc := &MockRootCoordClientInterface{} qc := getQueryCoordClient() + qc.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{}, nil).Maybe() qc.EXPECT().DescribeResourceGroup(mock.Anything, mock.Anything).Return(&querypb.DescribeResourceGroupResponse{ Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_UnexpectedError}, }, nil) @@ -3145,8 +3158,11 @@ func TestCreateCollectionTaskWithPartitionKey(t *testing.T) { err = task.Execute(ctx) assert.NoError(t, err) + qc := getQueryCoordClient() + qc.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{}, nil).Maybe() + // check default partitions - err = InitMetaCache(ctx, rc, nil, nil) + err = InitMetaCache(ctx, rc, qc, nil) assert.NoError(t, err) partitionNames, err := getDefaultPartitionsInPartitionKeyMode(ctx, "", task.CollectionName) assert.NoError(t, err) @@ -3225,6 +3241,7 @@ func TestPartitionKey(t *testing.T) { defer rc.Close() qc := getQueryCoordClient() + qc.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{}, nil).Maybe() ctx := context.Background() @@ -3480,6 +3497,7 @@ func TestClusteringKey(t *testing.T) { defer rc.Close() qc := getQueryCoordClient() + qc.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{}, nil).Maybe() ctx := context.Background() @@ -3662,6 +3680,7 @@ func TestTaskPartitionKeyIsolation(t *testing.T) { dc := NewDataCoordMock() defer dc.Close() qc := getQueryCoordClient() + qc.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{}, nil).Maybe() defer qc.Close() ctx := context.Background() mgr := newShardClientMgr()