diff --git a/Makefile b/Makefile index 72099e78c8..9f62206990 100644 --- a/Makefile +++ b/Makefile @@ -455,7 +455,8 @@ generate-mockery-utils: getdeps # tso.Allocator $(INSTALL_PATH)/mockery --name=Allocator --dir=internal/tso --output=internal/tso/mocks --filename=allocator.go --with-expecter --structname=Allocator --outpkg=mocktso $(INSTALL_PATH)/mockery --name=SessionInterface --dir=$(PWD)/internal/util/sessionutil --output=$(PWD)/internal/util/sessionutil --filename=mock_session.go --with-expecter --structname=MockSession --inpackage - + $(INSTALL_PATH)/mockery --name=GrpcClient --dir=$(PWD)/internal/util/grpcclient --output=$(PWD)/internal/mocks --filename=mock_grpc_client.go --with-expecter --structname=MockGrpcClient + generate-mockery-kv: getdeps $(INSTALL_PATH)/mockery --name=TxnKV --dir=$(PWD)/internal/kv --output=$(PWD)/internal/kv/mocks --filename=txn_kv.go --with-expecter $(INSTALL_PATH)/mockery --name=MetaKv --dir=$(PWD)/internal/kv --output=$(PWD)/internal/kv/mocks --filename=meta_kv.go --with-expecter diff --git a/internal/distributed/datacoord/client/client_test.go b/internal/distributed/datacoord/client/client_test.go index 775013301a..f051757e0b 100644 --- a/internal/distributed/datacoord/client/client_test.go +++ b/internal/distributed/datacoord/client/client_test.go @@ -24,15 +24,18 @@ import ( "testing" "time" - "github.com/cockroachdb/errors" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" "go.uber.org/zap" - "google.golang.org/grpc" + "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb" + "github.com/milvus-io/milvus/internal/mocks" "github.com/milvus-io/milvus/internal/proto/datapb" - "github.com/milvus-io/milvus/internal/util/mock" + "github.com/milvus-io/milvus/internal/proto/indexpb" + "github.com/milvus-io/milvus/internal/proto/internalpb" "github.com/milvus-io/milvus/pkg/log" "github.com/milvus-io/milvus/pkg/util/etcd" + "github.com/milvus-io/milvus/pkg/util/merr" "github.com/milvus-io/milvus/pkg/util/paramtable" ) @@ -61,186 +64,1762 @@ func Test_NewClient(t *testing.T) { assert.NoError(t, err) assert.NotNil(t, client) - checkFunc := func(retNotNil bool) { - retCheck := func(notNil bool, ret any, err error) { - if notNil { - assert.NotNil(t, ret) - assert.NoError(t, err) - } else { - assert.Nil(t, ret) - assert.Error(t, err) - } - } - - r1, err := client.GetComponentStates(ctx, nil) - retCheck(retNotNil, r1, err) - - r2, err := client.GetTimeTickChannel(ctx, nil) - retCheck(retNotNil, r2, err) - - r3, err := client.GetStatisticsChannel(ctx, nil) - retCheck(retNotNil, r3, err) - - r4, err := client.Flush(ctx, nil) - retCheck(retNotNil, r4, err) - - r5, err := client.AssignSegmentID(ctx, nil) - retCheck(retNotNil, r5, err) - - r6, err := client.GetSegmentInfo(ctx, nil) - retCheck(retNotNil, r6, err) - - r7, err := client.GetSegmentStates(ctx, nil) - retCheck(retNotNil, r7, err) - - r8, err := client.GetInsertBinlogPaths(ctx, nil) - retCheck(retNotNil, r8, err) - - r9, err := client.GetCollectionStatistics(ctx, nil) - retCheck(retNotNil, r9, err) - - r10, err := client.GetPartitionStatistics(ctx, nil) - retCheck(retNotNil, r10, err) - - r11, err := client.GetSegmentInfoChannel(ctx, nil) - retCheck(retNotNil, r11, err) - - // r12, err := client.SaveBinlogPaths(ctx, nil) - // retCheck(retNotNil, r12, err) - - r13, err := client.GetRecoveryInfo(ctx, nil) - retCheck(retNotNil, r13, err) - - r14, err := client.GetFlushedSegments(ctx, nil) - retCheck(retNotNil, r14, err) - - r15, err := client.GetMetrics(ctx, nil) - retCheck(retNotNil, r15, err) - - r17, err := client.GetCompactionState(ctx, nil) - retCheck(retNotNil, r17, err) - - r18, err := client.ManualCompaction(ctx, nil) - retCheck(retNotNil, r18, err) - - r19, err := client.GetCompactionStateWithPlans(ctx, nil) - retCheck(retNotNil, r19, err) - - r20, err := client.WatchChannels(ctx, nil) - retCheck(retNotNil, r20, err) - - r21, err := client.DropVirtualChannel(ctx, nil) - retCheck(retNotNil, r21, err) - - r22, err := client.SetSegmentState(ctx, nil) - retCheck(retNotNil, r22, err) - - r23, err := client.Import(ctx, nil) - retCheck(retNotNil, r23, err) - - r24, err := client.UpdateSegmentStatistics(ctx, nil) - retCheck(retNotNil, r24, err) - - r27, err := client.SaveImportSegment(ctx, nil) - retCheck(retNotNil, r27, err) - - r29, err := client.UnsetIsImportingState(ctx, nil) - retCheck(retNotNil, r29, err) - - r30, err := client.MarkSegmentsDropped(ctx, nil) - retCheck(retNotNil, r30, err) - - r31, err := client.ShowConfigurations(ctx, nil) - retCheck(retNotNil, r31, err) - - r32, err := client.CreateIndex(ctx, nil) - retCheck(retNotNil, r32, err) - - r33, err := client.DescribeIndex(ctx, nil) - retCheck(retNotNil, r33, err) - - r34, err := client.DropIndex(ctx, nil) - retCheck(retNotNil, r34, err) - - r35, err := client.GetIndexState(ctx, nil) - retCheck(retNotNil, r35, err) - - r36, err := client.GetIndexBuildProgress(ctx, nil) - retCheck(retNotNil, r36, err) - - r37, err := client.GetIndexInfos(ctx, nil) - retCheck(retNotNil, r37, err) - - r38, err := client.GetSegmentIndexState(ctx, nil) - retCheck(retNotNil, r38, err) - - r39, err := client.UpdateChannelCheckpoint(ctx, nil) - retCheck(retNotNil, r39, err) - - r, err := client.GetFlushAllState(ctx, nil) - retCheck(retNotNil, r, err) - - { - ret, err := client.BroadcastAlteredCollection(ctx, nil) - retCheck(retNotNil, ret, err) - } - - { - ret, err := client.CheckHealth(ctx, nil) - retCheck(retNotNil, ret, err) - } - - r40, err := client.GetRecoveryInfoV2(ctx, nil) - retCheck(retNotNil, r40, err) - - r41, err := client.GetIndexStatistics(ctx, nil) - retCheck(retNotNil, r41, err) - } - - client.grpcClient = &mock.GRPCClientBase[datapb.DataCoordClient]{ - GetGrpcClientErr: errors.New("dummy"), - } - - newFunc1 := func(cc *grpc.ClientConn) datapb.DataCoordClient { - return &mock.GrpcDataCoordClient{Err: nil} - } - client.grpcClient.SetNewGrpcClientFunc(newFunc1) - - checkFunc(false) - - // special case since this method didn't use recall() - ret, err := client.SaveBinlogPaths(ctx, nil) - assert.Nil(t, ret) - assert.Error(t, err) - - client.grpcClient = &mock.GRPCClientBase[datapb.DataCoordClient]{ - GetGrpcClientErr: nil, - } - newFunc2 := func(cc *grpc.ClientConn) datapb.DataCoordClient { - return &mock.GrpcDataCoordClient{Err: errors.New("dummy")} - } - client.grpcClient.SetNewGrpcClientFunc(newFunc2) - checkFunc(false) - - // special case since this method didn't use recall() - ret, err = client.SaveBinlogPaths(ctx, nil) - assert.Nil(t, ret) - assert.Error(t, err) - - client.grpcClient = &mock.GRPCClientBase[datapb.DataCoordClient]{ - GetGrpcClientErr: nil, - } - newFunc3 := func(cc *grpc.ClientConn) datapb.DataCoordClient { - return &mock.GrpcDataCoordClient{Err: nil} - } - client.grpcClient.SetNewGrpcClientFunc(newFunc3) - checkFunc(true) - - // special case since this method didn't use recall() - ret, err = client.SaveBinlogPaths(ctx, nil) - assert.NotNil(t, ret) - assert.NoError(t, err) - err = client.Close() assert.NoError(t, err) } + +func Test_GetComponentStates(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().GetComponentStates(mock.Anything, mock.Anything).Return(&milvuspb.ComponentStates{ + Status: merr.Success(), + }, nil) + _, err = client.GetComponentStates(ctx, &milvuspb.GetComponentStatesRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().GetComponentStates(mock.Anything, mock.Anything).Return(&milvuspb.ComponentStates{ + Status: merr.Status(merr.ErrServiceNotReady), + }, nil) + + _, err = client.GetComponentStates(ctx, &milvuspb.GetComponentStatesRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.GetComponentStates(ctx, &milvuspb.GetComponentStatesRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_GetTimeTickChannel(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().GetTimeTickChannel(mock.Anything, mock.Anything).Return(&milvuspb.StringResponse{ + Status: merr.Success(), + }, nil) + _, err = client.GetTimeTickChannel(ctx, &internalpb.GetTimeTickChannelRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().GetTimeTickChannel(mock.Anything, mock.Anything).Return(&milvuspb.StringResponse{ + Status: merr.Status(merr.ErrServiceNotReady), + }, nil) + + _, err = client.GetTimeTickChannel(ctx, &internalpb.GetTimeTickChannelRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.GetTimeTickChannel(ctx, &internalpb.GetTimeTickChannelRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_GetStatisticsChannel(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().GetStatisticsChannel(mock.Anything, mock.Anything).Return(&milvuspb.StringResponse{ + Status: merr.Success(), + }, nil) + _, err = client.GetStatisticsChannel(ctx, &internalpb.GetStatisticsChannelRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().GetStatisticsChannel(mock.Anything, mock.Anything).Return(&milvuspb.StringResponse{ + Status: merr.Status(merr.ErrServiceNotReady), + }, nil) + + _, err = client.GetStatisticsChannel(ctx, &internalpb.GetStatisticsChannelRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.GetStatisticsChannel(ctx, &internalpb.GetStatisticsChannelRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_Flush(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().GetNodeID().Return(1) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().Flush(mock.Anything, mock.Anything).Return(&datapb.FlushResponse{ + Status: merr.Success(), + }, nil) + _, err = client.Flush(ctx, &datapb.FlushRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().Flush(mock.Anything, mock.Anything).Return(&datapb.FlushResponse{ + Status: merr.Status(merr.ErrServiceNotReady), + }, nil) + + _, err = client.Flush(ctx, &datapb.FlushRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.Flush(ctx, &datapb.FlushRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_AssignSegmentID(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().AssignSegmentID(mock.Anything, mock.Anything).Return(&datapb.AssignSegmentIDResponse{ + Status: merr.Success(), + }, nil) + _, err = client.AssignSegmentID(ctx, &datapb.AssignSegmentIDRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().AssignSegmentID(mock.Anything, mock.Anything).Return(&datapb.AssignSegmentIDResponse{ + Status: merr.Status(merr.ErrServiceNotReady), + }, nil) + + _, err = client.AssignSegmentID(ctx, &datapb.AssignSegmentIDRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.AssignSegmentID(ctx, &datapb.AssignSegmentIDRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_GetSegmentStates(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().GetNodeID().Return(1) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().GetSegmentStates(mock.Anything, mock.Anything).Return(&datapb.GetSegmentStatesResponse{ + Status: merr.Success(), + }, nil) + _, err = client.GetSegmentStates(ctx, &datapb.GetSegmentStatesRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().GetSegmentStates(mock.Anything, mock.Anything).Return(&datapb.GetSegmentStatesResponse{ + Status: merr.Status(merr.ErrServiceNotReady), + }, nil) + + _, err = client.GetSegmentStates(ctx, &datapb.GetSegmentStatesRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.GetSegmentStates(ctx, &datapb.GetSegmentStatesRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_GetInsertBinlogPaths(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().GetNodeID().Return(1) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().GetInsertBinlogPaths(mock.Anything, mock.Anything).Return(&datapb.GetInsertBinlogPathsResponse{ + Status: merr.Success(), + }, nil) + _, err = client.GetInsertBinlogPaths(ctx, &datapb.GetInsertBinlogPathsRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().GetInsertBinlogPaths(mock.Anything, mock.Anything).Return(&datapb.GetInsertBinlogPathsResponse{ + Status: merr.Status(merr.ErrServiceNotReady), + }, nil) + + _, err = client.GetInsertBinlogPaths(ctx, &datapb.GetInsertBinlogPathsRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.GetInsertBinlogPaths(ctx, &datapb.GetInsertBinlogPathsRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_GetCollectionStatistics(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().GetNodeID().Return(1) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().GetCollectionStatistics(mock.Anything, mock.Anything).Return(&datapb.GetCollectionStatisticsResponse{ + Status: merr.Success(), + }, nil) + _, err = client.GetCollectionStatistics(ctx, &datapb.GetCollectionStatisticsRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().GetCollectionStatistics(mock.Anything, mock.Anything).Return(&datapb.GetCollectionStatisticsResponse{ + Status: merr.Status(merr.ErrServiceNotReady), + }, nil) + + _, err = client.GetCollectionStatistics(ctx, &datapb.GetCollectionStatisticsRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.GetCollectionStatistics(ctx, &datapb.GetCollectionStatisticsRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_GetPartitionStatistics(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().GetNodeID().Return(1) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().GetPartitionStatistics(mock.Anything, mock.Anything).Return(&datapb.GetPartitionStatisticsResponse{ + Status: merr.Success(), + }, nil) + _, err = client.GetPartitionStatistics(ctx, &datapb.GetPartitionStatisticsRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().GetPartitionStatistics(mock.Anything, mock.Anything).Return(&datapb.GetPartitionStatisticsResponse{ + Status: merr.Status(merr.ErrServiceNotReady), + }, nil) + + _, err = client.GetPartitionStatistics(ctx, &datapb.GetPartitionStatisticsRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.GetPartitionStatistics(ctx, &datapb.GetPartitionStatisticsRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_GetSegmentInfoChannel(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().GetSegmentInfoChannel(mock.Anything, mock.Anything).Return(&milvuspb.StringResponse{ + Status: merr.Success(), + }, nil) + _, err = client.GetSegmentInfoChannel(ctx, &datapb.GetSegmentInfoChannelRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().GetSegmentInfoChannel(mock.Anything, mock.Anything).Return(&milvuspb.StringResponse{ + Status: merr.Status(merr.ErrServiceNotReady), + }, nil) + + _, err = client.GetSegmentInfoChannel(ctx, &datapb.GetSegmentInfoChannelRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.GetSegmentInfoChannel(ctx, &datapb.GetSegmentInfoChannelRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_GetSegmentInfo(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().GetNodeID().Return(1) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().GetSegmentInfo(mock.Anything, mock.Anything).Return(&datapb.GetSegmentInfoResponse{ + Status: merr.Success(), + }, nil) + _, err = client.GetSegmentInfo(ctx, &datapb.GetSegmentInfoRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().GetSegmentInfo(mock.Anything, mock.Anything).Return(&datapb.GetSegmentInfoResponse{ + Status: merr.Status(merr.ErrServiceNotReady), + }, nil) + + _, err = client.GetSegmentInfo(ctx, &datapb.GetSegmentInfoRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.GetSegmentInfo(ctx, &datapb.GetSegmentInfoRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_SaveBinlogPaths(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().GetNodeID().Return(1) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().GetSegmentInfo(mock.Anything, mock.Anything).Return(&datapb.GetSegmentInfoResponse{ + Status: merr.Success(), + }, nil) + _, err = client.GetSegmentInfo(ctx, &datapb.GetSegmentInfoRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().GetSegmentInfo(mock.Anything, mock.Anything).Return(&datapb.GetSegmentInfoResponse{ + Status: merr.Status(merr.ErrServiceNotReady), + }, nil) + + _, err = client.GetSegmentInfo(ctx, &datapb.GetSegmentInfoRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.GetSegmentInfo(ctx, &datapb.GetSegmentInfoRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_GetRecoveryInfo(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().GetNodeID().Return(1) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().GetRecoveryInfo(mock.Anything, mock.Anything).Return(&datapb.GetRecoveryInfoResponse{ + Status: merr.Success(), + }, nil) + _, err = client.GetRecoveryInfo(ctx, &datapb.GetRecoveryInfoRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().GetRecoveryInfo(mock.Anything, mock.Anything).Return(&datapb.GetRecoveryInfoResponse{ + Status: merr.Status(merr.ErrServiceNotReady), + }, nil) + + _, err = client.GetRecoveryInfo(ctx, &datapb.GetRecoveryInfoRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.GetRecoveryInfo(ctx, &datapb.GetRecoveryInfoRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_GetRecoveryInfoV2(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().GetRecoveryInfoV2(mock.Anything, mock.Anything).Return(&datapb.GetRecoveryInfoResponseV2{ + Status: merr.Success(), + }, nil) + _, err = client.GetRecoveryInfoV2(ctx, &datapb.GetRecoveryInfoRequestV2{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().GetRecoveryInfoV2(mock.Anything, mock.Anything).Return(&datapb.GetRecoveryInfoResponseV2{ + Status: merr.Status(merr.ErrServiceNotReady), + }, nil) + + _, err = client.GetRecoveryInfoV2(ctx, &datapb.GetRecoveryInfoRequestV2{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.GetRecoveryInfoV2(ctx, &datapb.GetRecoveryInfoRequestV2{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_GetFlushedSegments(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().GetNodeID().Return(1) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().GetFlushedSegments(mock.Anything, mock.Anything).Return(&datapb.GetFlushedSegmentsResponse{ + Status: merr.Success(), + }, nil) + _, err = client.GetFlushedSegments(ctx, &datapb.GetFlushedSegmentsRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().GetFlushedSegments(mock.Anything, mock.Anything).Return(&datapb.GetFlushedSegmentsResponse{ + Status: merr.Status(merr.ErrServiceNotReady), + }, nil) + + _, err = client.GetFlushedSegments(ctx, &datapb.GetFlushedSegmentsRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.GetFlushedSegments(ctx, &datapb.GetFlushedSegmentsRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_GetSegmentsByStates(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().GetNodeID().Return(1) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().GetSegmentsByStates(mock.Anything, mock.Anything).Return(&datapb.GetSegmentsByStatesResponse{ + Status: merr.Success(), + }, nil) + _, err = client.GetSegmentsByStates(ctx, &datapb.GetSegmentsByStatesRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().GetSegmentsByStates(mock.Anything, mock.Anything).Return(&datapb.GetSegmentsByStatesResponse{ + Status: merr.Status(merr.ErrServiceNotReady), + }, nil) + + _, err = client.GetSegmentsByStates(ctx, &datapb.GetSegmentsByStatesRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.GetSegmentsByStates(ctx, &datapb.GetSegmentsByStatesRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_ShowConfigurations(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().GetNodeID().Return(1) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().ShowConfigurations(mock.Anything, mock.Anything).Return(&internalpb.ShowConfigurationsResponse{ + Status: merr.Success(), + }, nil) + _, err = client.ShowConfigurations(ctx, &internalpb.ShowConfigurationsRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().ShowConfigurations(mock.Anything, mock.Anything).Return(&internalpb.ShowConfigurationsResponse{ + Status: merr.Status(merr.ErrServiceNotReady), + }, nil) + + _, err = client.ShowConfigurations(ctx, &internalpb.ShowConfigurationsRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.ShowConfigurations(ctx, &internalpb.ShowConfigurationsRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_GetMetrics(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().GetNodeID().Return(1) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(&milvuspb.GetMetricsResponse{ + Status: merr.Success(), + }, nil) + _, err = client.GetMetrics(ctx, &milvuspb.GetMetricsRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(&milvuspb.GetMetricsResponse{ + Status: merr.Status(merr.ErrServiceNotReady), + }, nil) + + _, err = client.GetMetrics(ctx, &milvuspb.GetMetricsRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.GetMetrics(ctx, &milvuspb.GetMetricsRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_ManualCompaction(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().ManualCompaction(mock.Anything, mock.Anything).Return(&milvuspb.ManualCompactionResponse{ + Status: merr.Success(), + }, nil) + _, err = client.ManualCompaction(ctx, &milvuspb.ManualCompactionRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().ManualCompaction(mock.Anything, mock.Anything).Return(&milvuspb.ManualCompactionResponse{ + Status: merr.Status(merr.ErrServiceNotReady), + }, nil) + + _, err = client.ManualCompaction(ctx, &milvuspb.ManualCompactionRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.ManualCompaction(ctx, &milvuspb.ManualCompactionRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_GetCompactionState(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().GetCompactionState(mock.Anything, mock.Anything).Return(&milvuspb.GetCompactionStateResponse{ + Status: merr.Success(), + }, nil) + _, err = client.GetCompactionState(ctx, &milvuspb.GetCompactionStateRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().GetCompactionState(mock.Anything, mock.Anything).Return(&milvuspb.GetCompactionStateResponse{ + Status: merr.Status(merr.ErrServiceNotReady), + }, nil) + + _, err = client.GetCompactionState(ctx, &milvuspb.GetCompactionStateRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.GetCompactionState(ctx, &milvuspb.GetCompactionStateRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_GetCompactionStateWithPlans(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().GetCompactionStateWithPlans(mock.Anything, mock.Anything).Return(&milvuspb.GetCompactionPlansResponse{ + Status: merr.Success(), + }, nil) + _, err = client.GetCompactionStateWithPlans(ctx, &milvuspb.GetCompactionPlansRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().GetCompactionStateWithPlans(mock.Anything, mock.Anything).Return(&milvuspb.GetCompactionPlansResponse{ + Status: merr.Status(merr.ErrServiceNotReady), + }, nil) + + _, err = client.GetCompactionStateWithPlans(ctx, &milvuspb.GetCompactionPlansRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.GetCompactionStateWithPlans(ctx, &milvuspb.GetCompactionPlansRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_WatchChannels(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().WatchChannels(mock.Anything, mock.Anything).Return(&datapb.WatchChannelsResponse{ + Status: merr.Success(), + }, nil) + _, err = client.WatchChannels(ctx, &datapb.WatchChannelsRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().WatchChannels(mock.Anything, mock.Anything).Return(&datapb.WatchChannelsResponse{ + Status: merr.Status(merr.ErrServiceNotReady), + }, nil) + + _, err = client.WatchChannels(ctx, &datapb.WatchChannelsRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.WatchChannels(ctx, &datapb.WatchChannelsRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_GetFlushState(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().GetFlushState(mock.Anything, mock.Anything).Return(&milvuspb.GetFlushStateResponse{ + Status: merr.Success(), + }, nil) + _, err = client.GetFlushState(ctx, &datapb.GetFlushStateRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().GetFlushState(mock.Anything, mock.Anything).Return(&milvuspb.GetFlushStateResponse{ + Status: merr.Status(merr.ErrServiceNotReady), + }, nil) + + _, err = client.GetFlushState(ctx, &datapb.GetFlushStateRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.GetFlushState(ctx, &datapb.GetFlushStateRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_GetFlushAllState(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().GetFlushAllState(mock.Anything, mock.Anything).Return(&milvuspb.GetFlushAllStateResponse{ + Status: merr.Success(), + }, nil) + _, err = client.GetFlushAllState(ctx, &milvuspb.GetFlushAllStateRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().GetFlushAllState(mock.Anything, mock.Anything).Return(&milvuspb.GetFlushAllStateResponse{ + Status: merr.Status(merr.ErrServiceNotReady), + }, nil) + + _, err = client.GetFlushAllState(ctx, &milvuspb.GetFlushAllStateRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.GetFlushAllState(ctx, &milvuspb.GetFlushAllStateRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_DropVirtualChannel(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().GetNodeID().Return(1) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().DropVirtualChannel(mock.Anything, mock.Anything).Return(&datapb.DropVirtualChannelResponse{ + Status: merr.Success(), + }, nil) + _, err = client.DropVirtualChannel(ctx, &datapb.DropVirtualChannelRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().DropVirtualChannel(mock.Anything, mock.Anything).Return(&datapb.DropVirtualChannelResponse{ + Status: merr.Status(merr.ErrServiceNotReady), + }, nil) + + _, err = client.DropVirtualChannel(ctx, &datapb.DropVirtualChannelRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.DropVirtualChannel(ctx, &datapb.DropVirtualChannelRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_SetSegmentState(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().GetNodeID().Return(1) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().SetSegmentState(mock.Anything, mock.Anything).Return(&datapb.SetSegmentStateResponse{ + Status: merr.Success(), + }, nil) + _, err = client.SetSegmentState(ctx, &datapb.SetSegmentStateRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().SetSegmentState(mock.Anything, mock.Anything).Return(&datapb.SetSegmentStateResponse{ + Status: merr.Status(merr.ErrServiceNotReady), + }, nil) + + _, err = client.SetSegmentState(ctx, &datapb.SetSegmentStateRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.SetSegmentState(ctx, &datapb.SetSegmentStateRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_Import(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().GetNodeID().Return(1) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().Import(mock.Anything, mock.Anything).Return(&datapb.ImportTaskResponse{ + Status: merr.Success(), + }, nil) + _, err = client.Import(ctx, &datapb.ImportTaskRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().Import(mock.Anything, mock.Anything).Return(&datapb.ImportTaskResponse{ + Status: merr.Status(merr.ErrServiceNotReady), + }, nil) + + _, err = client.Import(ctx, &datapb.ImportTaskRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.Import(ctx, &datapb.ImportTaskRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_UpdateSegmentStatistics(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().GetNodeID().Return(1) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().UpdateSegmentStatistics(mock.Anything, mock.Anything).Return(merr.Success(), nil) + _, err = client.UpdateSegmentStatistics(ctx, &datapb.UpdateSegmentStatisticsRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().UpdateSegmentStatistics(mock.Anything, mock.Anything).Return(merr.Status(merr.ErrServiceNotReady), nil) + + _, err = client.UpdateSegmentStatistics(ctx, &datapb.UpdateSegmentStatisticsRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.UpdateSegmentStatistics(ctx, &datapb.UpdateSegmentStatisticsRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_UpdateChannelCheckpoint(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().GetNodeID().Return(1) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().UpdateChannelCheckpoint(mock.Anything, mock.Anything).Return(merr.Success(), nil) + _, err = client.UpdateChannelCheckpoint(ctx, &datapb.UpdateChannelCheckpointRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().UpdateChannelCheckpoint(mock.Anything, mock.Anything).Return(merr.Status(merr.ErrServiceNotReady), nil) + + _, err = client.UpdateChannelCheckpoint(ctx, &datapb.UpdateChannelCheckpointRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.UpdateChannelCheckpoint(ctx, &datapb.UpdateChannelCheckpointRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_SaveImportSegment(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().GetNodeID().Return(1) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().SaveImportSegment(mock.Anything, mock.Anything).Return(merr.Success(), nil) + _, err = client.SaveImportSegment(ctx, &datapb.SaveImportSegmentRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().SaveImportSegment(mock.Anything, mock.Anything).Return(merr.Status(merr.ErrServiceNotReady), nil) + + _, err = client.SaveImportSegment(ctx, &datapb.SaveImportSegmentRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.SaveImportSegment(ctx, &datapb.SaveImportSegmentRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_UnsetIsImportingState(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().GetNodeID().Return(1) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().UnsetIsImportingState(mock.Anything, mock.Anything).Return(merr.Success(), nil) + _, err = client.UnsetIsImportingState(ctx, &datapb.UnsetIsImportingStateRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().UnsetIsImportingState(mock.Anything, mock.Anything).Return(merr.Status(merr.ErrServiceNotReady), nil) + + _, err = client.UnsetIsImportingState(ctx, &datapb.UnsetIsImportingStateRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.UnsetIsImportingState(ctx, &datapb.UnsetIsImportingStateRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_MarkSegmentsDropped(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().GetNodeID().Return(1) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().MarkSegmentsDropped(mock.Anything, mock.Anything).Return(merr.Success(), nil) + _, err = client.MarkSegmentsDropped(ctx, &datapb.MarkSegmentsDroppedRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().MarkSegmentsDropped(mock.Anything, mock.Anything).Return(merr.Status(merr.ErrServiceNotReady), nil) + + _, err = client.MarkSegmentsDropped(ctx, &datapb.MarkSegmentsDroppedRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.MarkSegmentsDropped(ctx, &datapb.MarkSegmentsDroppedRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_BroadcastAlteredCollection(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().BroadcastAlteredCollection(mock.Anything, mock.Anything).Return(merr.Success(), nil) + _, err = client.BroadcastAlteredCollection(ctx, &datapb.AlterCollectionRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().BroadcastAlteredCollection(mock.Anything, mock.Anything).Return(merr.Status(merr.ErrServiceNotReady), nil) + + _, err = client.BroadcastAlteredCollection(ctx, &datapb.AlterCollectionRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.BroadcastAlteredCollection(ctx, &datapb.AlterCollectionRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_CheckHealth(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().CheckHealth(mock.Anything, mock.Anything).Return(&milvuspb.CheckHealthResponse{Status: merr.Success()}, nil) + _, err = client.CheckHealth(ctx, &milvuspb.CheckHealthRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().CheckHealth(mock.Anything, mock.Anything).Return(&milvuspb.CheckHealthResponse{Status: merr.Status(merr.ErrServiceNotReady)}, nil) + + _, err = client.CheckHealth(ctx, &milvuspb.CheckHealthRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.CheckHealth(ctx, &milvuspb.CheckHealthRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_GcConfirm(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().GcConfirm(mock.Anything, mock.Anything).Return(&datapb.GcConfirmResponse{Status: merr.Success()}, nil) + _, err = client.GcConfirm(ctx, &datapb.GcConfirmRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().GcConfirm(mock.Anything, mock.Anything).Return(&datapb.GcConfirmResponse{Status: merr.Status(merr.ErrServiceNotReady)}, nil) + + _, err = client.GcConfirm(ctx, &datapb.GcConfirmRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.GcConfirm(ctx, &datapb.GcConfirmRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_CreateIndex(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().CreateIndex(mock.Anything, mock.Anything).Return(merr.Success(), nil) + _, err = client.CreateIndex(ctx, &indexpb.CreateIndexRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().CreateIndex(mock.Anything, mock.Anything).Return(merr.Status(merr.ErrServiceNotReady), nil) + + _, err = client.CreateIndex(ctx, &indexpb.CreateIndexRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.CreateIndex(ctx, &indexpb.CreateIndexRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_GetSegmentIndexState(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().GetSegmentIndexState(mock.Anything, mock.Anything).Return(&indexpb.GetSegmentIndexStateResponse{Status: merr.Success()}, nil) + _, err = client.GetSegmentIndexState(ctx, &indexpb.GetSegmentIndexStateRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().GetSegmentIndexState(mock.Anything, mock.Anything).Return(&indexpb.GetSegmentIndexStateResponse{Status: merr.Status(err)}, nil) + + _, err = client.GetSegmentIndexState(ctx, &indexpb.GetSegmentIndexStateRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.GetSegmentIndexState(ctx, &indexpb.GetSegmentIndexStateRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_GetIndexState(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().GetIndexState(mock.Anything, mock.Anything).Return(&indexpb.GetIndexStateResponse{Status: merr.Success()}, nil) + _, err = client.GetIndexState(ctx, &indexpb.GetIndexStateRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().GetIndexState(mock.Anything, mock.Anything).Return(&indexpb.GetIndexStateResponse{Status: merr.Status(err)}, nil) + + _, err = client.GetIndexState(ctx, &indexpb.GetIndexStateRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.GetIndexState(ctx, &indexpb.GetIndexStateRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_GetIndexInfos(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().GetIndexInfos(mock.Anything, mock.Anything).Return(&indexpb.GetIndexInfoResponse{Status: merr.Success()}, nil) + _, err = client.GetIndexInfos(ctx, &indexpb.GetIndexInfoRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().GetIndexInfos(mock.Anything, mock.Anything).Return(&indexpb.GetIndexInfoResponse{Status: merr.Status(err)}, nil) + + _, err = client.GetIndexInfos(ctx, &indexpb.GetIndexInfoRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.GetIndexInfos(ctx, &indexpb.GetIndexInfoRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_DescribeIndex(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().DescribeIndex(mock.Anything, mock.Anything).Return(&indexpb.DescribeIndexResponse{Status: merr.Success()}, nil) + _, err = client.DescribeIndex(ctx, &indexpb.DescribeIndexRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().DescribeIndex(mock.Anything, mock.Anything).Return(&indexpb.DescribeIndexResponse{Status: merr.Status(err)}, nil) + + _, err = client.DescribeIndex(ctx, &indexpb.DescribeIndexRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.DescribeIndex(ctx, &indexpb.DescribeIndexRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_GetIndexStatistics(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().GetIndexStatistics(mock.Anything, mock.Anything).Return(&indexpb.GetIndexStatisticsResponse{Status: merr.Success()}, nil) + _, err = client.GetIndexStatistics(ctx, &indexpb.GetIndexStatisticsRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().GetIndexStatistics(mock.Anything, mock.Anything).Return(&indexpb.GetIndexStatisticsResponse{Status: merr.Status(err)}, nil) + + _, err = client.GetIndexStatistics(ctx, &indexpb.GetIndexStatisticsRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.GetIndexStatistics(ctx, &indexpb.GetIndexStatisticsRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_GetIndexBuildProgress(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().GetIndexBuildProgress(mock.Anything, mock.Anything).Return(&indexpb.GetIndexBuildProgressResponse{Status: merr.Success()}, nil) + _, err = client.GetIndexBuildProgress(ctx, &indexpb.GetIndexBuildProgressRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().GetIndexBuildProgress(mock.Anything, mock.Anything).Return(&indexpb.GetIndexBuildProgressResponse{Status: merr.Status(err)}, nil) + + _, err = client.GetIndexBuildProgress(ctx, &indexpb.GetIndexBuildProgressRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.GetIndexBuildProgress(ctx, &indexpb.GetIndexBuildProgressRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_DropIndex(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().DropIndex(mock.Anything, mock.Anything).Return(merr.Success(), nil) + _, err = client.DropIndex(ctx, &indexpb.DropIndexRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().DropIndex(mock.Anything, mock.Anything).Return(merr.Status(err), nil) + + _, err = client.DropIndex(ctx, &indexpb.DropIndexRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.DropIndex(ctx, &indexpb.DropIndexRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_ReportDataNodeTtMsgs(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockDataCoordClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[datapb.DataCoordClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(datapb.DataCoordClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().ReportDataNodeTtMsgs(mock.Anything, mock.Anything).Return(merr.Success(), nil) + _, err = client.ReportDataNodeTtMsgs(ctx, &datapb.ReportDataNodeTtMsgsRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().ReportDataNodeTtMsgs(mock.Anything, mock.Anything).Return(merr.Status(err), nil) + + _, err = client.ReportDataNodeTtMsgs(ctx, &datapb.ReportDataNodeTtMsgsRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.ReportDataNodeTtMsgs(ctx, &datapb.ReportDataNodeTtMsgsRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} diff --git a/internal/distributed/datacoord/service_test.go b/internal/distributed/datacoord/service_test.go index cf8429bfd1..1532201808 100644 --- a/internal/distributed/datacoord/service_test.go +++ b/internal/distributed/datacoord/service_test.go @@ -18,673 +18,319 @@ package grpcdatacoord import ( "context" - "fmt" "testing" "github.com/cockroachdb/errors" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" "github.com/tikv/client-go/v2/txnkv" - clientv3 "go.etcd.io/etcd/client/v3" - "github.com/milvus-io/milvus-proto/go-api/v2/commonpb" "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb" + "github.com/milvus-io/milvus/internal/mocks" "github.com/milvus-io/milvus/internal/proto/datapb" "github.com/milvus-io/milvus/internal/proto/indexpb" "github.com/milvus-io/milvus/internal/proto/internalpb" - "github.com/milvus-io/milvus/internal/types" "github.com/milvus-io/milvus/pkg/util/merr" "github.com/milvus-io/milvus/pkg/util/paramtable" "github.com/milvus-io/milvus/pkg/util/tikv" ) -type MockDataCoord struct { - types.DataCoord - - states *milvuspb.ComponentStates - status *commonpb.Status - err error - initErr error - startErr error - stopErr error - regErr error - strResp *milvuspb.StringResponse - infoResp *datapb.GetSegmentInfoResponse - flushResp *datapb.FlushResponse - assignResp *datapb.AssignSegmentIDResponse - segStateResp *datapb.GetSegmentStatesResponse - binResp *datapb.GetInsertBinlogPathsResponse - colStatResp *datapb.GetCollectionStatisticsResponse - partStatResp *datapb.GetPartitionStatisticsResponse - recoverResp *datapb.GetRecoveryInfoResponse - flushSegResp *datapb.GetFlushedSegmentsResponse - SegByStatesResp *datapb.GetSegmentsByStatesResponse - configResp *internalpb.ShowConfigurationsResponse - metricResp *milvuspb.GetMetricsResponse - compactionStateResp *milvuspb.GetCompactionStateResponse - manualCompactionResp *milvuspb.ManualCompactionResponse - compactionPlansResp *milvuspb.GetCompactionPlansResponse - watchChannelsResp *datapb.WatchChannelsResponse - getFlushStateResp *milvuspb.GetFlushStateResponse - getFlushAllStateResp *milvuspb.GetFlushAllStateResponse - dropVChanResp *datapb.DropVirtualChannelResponse - setSegmentStateResp *datapb.SetSegmentStateResponse - importResp *datapb.ImportTaskResponse - updateSegStatResp *commonpb.Status - updateChanPos *commonpb.Status - addSegmentResp *commonpb.Status - unsetIsImportingStateResp *commonpb.Status - markSegmentsDroppedResp *commonpb.Status - broadCastResp *commonpb.Status - - createIndexResp *commonpb.Status - describeIndexResp *indexpb.DescribeIndexResponse - getIndexStatisticsResp *indexpb.GetIndexStatisticsResponse - dropIndexResp *commonpb.Status - getIndexStateResp *indexpb.GetIndexStateResponse - getIndexBuildProgressResp *indexpb.GetIndexBuildProgressResponse - getSegmentIndexStateResp *indexpb.GetSegmentIndexStateResponse - getIndexInfosResp *indexpb.GetIndexInfoResponse -} - -func (m *MockDataCoord) Init() error { - return m.initErr -} - -func (m *MockDataCoord) Start() error { - return m.startErr -} - -func (m *MockDataCoord) Stop() error { - return m.stopErr -} - -func (m *MockDataCoord) Register() error { - return m.regErr -} - -func (*MockDataCoord) SetAddress(address string) { -} - -func (m *MockDataCoord) SetEtcdClient(etcdClient *clientv3.Client) { -} - -func (m *MockDataCoord) SetTiKVClient(client *txnkv.Client) { -} - -func (m *MockDataCoord) SetRootCoordClient(rootCoord types.RootCoordClient) { -} - -func (m *MockDataCoord) SetDataNodeCreator(func(context.Context, string, int64) (types.DataNodeClient, error)) { -} - -func (m *MockDataCoord) SetIndexNodeCreator(func(context.Context, string, int64) (types.IndexNodeClient, error)) { -} - -func (m *MockDataCoord) GetComponentStates(ctx context.Context, req *milvuspb.GetComponentStatesRequest) (*milvuspb.ComponentStates, error) { - return m.states, m.err -} - -func (m *MockDataCoord) GetTimeTickChannel(ctx context.Context, req *internalpb.GetTimeTickChannelRequest) (*milvuspb.StringResponse, error) { - return m.strResp, m.err -} - -func (m *MockDataCoord) GetStatisticsChannel(ctx context.Context, req *internalpb.GetStatisticsChannelRequest) (*milvuspb.StringResponse, error) { - return m.strResp, m.err -} - -func (m *MockDataCoord) GetSegmentInfo(ctx context.Context, req *datapb.GetSegmentInfoRequest) (*datapb.GetSegmentInfoResponse, error) { - return m.infoResp, m.err -} - -func (m *MockDataCoord) Flush(ctx context.Context, req *datapb.FlushRequest) (*datapb.FlushResponse, error) { - return m.flushResp, m.err -} - -func (m *MockDataCoord) AssignSegmentID(ctx context.Context, req *datapb.AssignSegmentIDRequest) (*datapb.AssignSegmentIDResponse, error) { - return m.assignResp, m.err -} - -func (m *MockDataCoord) GetSegmentStates(ctx context.Context, req *datapb.GetSegmentStatesRequest) (*datapb.GetSegmentStatesResponse, error) { - return m.segStateResp, m.err -} - -func (m *MockDataCoord) GetInsertBinlogPaths(ctx context.Context, req *datapb.GetInsertBinlogPathsRequest) (*datapb.GetInsertBinlogPathsResponse, error) { - return m.binResp, m.err -} - -func (m *MockDataCoord) GetCollectionStatistics(ctx context.Context, req *datapb.GetCollectionStatisticsRequest) (*datapb.GetCollectionStatisticsResponse, error) { - return m.colStatResp, m.err -} - -func (m *MockDataCoord) GetPartitionStatistics(ctx context.Context, req *datapb.GetPartitionStatisticsRequest) (*datapb.GetPartitionStatisticsResponse, error) { - return m.partStatResp, m.err -} - -func (m *MockDataCoord) GetSegmentInfoChannel(ctx context.Context, req *datapb.GetSegmentInfoChannelRequest) (*milvuspb.StringResponse, error) { - return m.strResp, m.err -} - -func (m *MockDataCoord) SaveBinlogPaths(ctx context.Context, req *datapb.SaveBinlogPathsRequest) (*commonpb.Status, error) { - return m.status, m.err -} - -func (m *MockDataCoord) GetRecoveryInfo(ctx context.Context, req *datapb.GetRecoveryInfoRequest) (*datapb.GetRecoveryInfoResponse, error) { - return m.recoverResp, m.err -} - -func (m *MockDataCoord) GetFlushedSegments(ctx context.Context, req *datapb.GetFlushedSegmentsRequest) (*datapb.GetFlushedSegmentsResponse, error) { - return m.flushSegResp, m.err -} - -func (m *MockDataCoord) GetSegmentsByStates(ctx context.Context, req *datapb.GetSegmentsByStatesRequest) (*datapb.GetSegmentsByStatesResponse, error) { - return m.SegByStatesResp, m.err -} - -func (m *MockDataCoord) ShowConfigurations(ctx context.Context, req *internalpb.ShowConfigurationsRequest) (*internalpb.ShowConfigurationsResponse, error) { - return m.configResp, m.err -} - -func (m *MockDataCoord) GetMetrics(ctx context.Context, req *milvuspb.GetMetricsRequest) (*milvuspb.GetMetricsResponse, error) { - return m.metricResp, m.err -} - -func (m *MockDataCoord) CompleteCompaction(ctx context.Context, req *datapb.CompactionPlanResult) (*commonpb.Status, error) { - return m.status, m.err -} - -func (m *MockDataCoord) ManualCompaction(ctx context.Context, req *milvuspb.ManualCompactionRequest) (*milvuspb.ManualCompactionResponse, error) { - return m.manualCompactionResp, m.err -} - -func (m *MockDataCoord) GetCompactionState(ctx context.Context, req *milvuspb.GetCompactionStateRequest) (*milvuspb.GetCompactionStateResponse, error) { - return m.compactionStateResp, m.err -} - -func (m *MockDataCoord) GetCompactionStateWithPlans(ctx context.Context, req *milvuspb.GetCompactionPlansRequest) (*milvuspb.GetCompactionPlansResponse, error) { - return m.compactionPlansResp, m.err -} - -func (m *MockDataCoord) WatchChannels(ctx context.Context, req *datapb.WatchChannelsRequest) (*datapb.WatchChannelsResponse, error) { - return m.watchChannelsResp, m.err -} - -func (m *MockDataCoord) GetFlushState(ctx context.Context, req *datapb.GetFlushStateRequest) (*milvuspb.GetFlushStateResponse, error) { - return m.getFlushStateResp, m.err -} - -func (m *MockDataCoord) GetFlushAllState(ctx context.Context, req *milvuspb.GetFlushAllStateRequest) (*milvuspb.GetFlushAllStateResponse, error) { - return m.getFlushAllStateResp, m.err -} - -func (m *MockDataCoord) DropVirtualChannel(ctx context.Context, req *datapb.DropVirtualChannelRequest) (*datapb.DropVirtualChannelResponse, error) { - return m.dropVChanResp, m.err -} - -func (m *MockDataCoord) SetSegmentState(ctx context.Context, req *datapb.SetSegmentStateRequest) (*datapb.SetSegmentStateResponse, error) { - return m.setSegmentStateResp, m.err -} - -func (m *MockDataCoord) Import(ctx context.Context, req *datapb.ImportTaskRequest) (*datapb.ImportTaskResponse, error) { - return m.importResp, m.err -} - -func (m *MockDataCoord) UpdateSegmentStatistics(ctx context.Context, req *datapb.UpdateSegmentStatisticsRequest) (*commonpb.Status, error) { - return m.updateSegStatResp, m.err -} - -func (m *MockDataCoord) UpdateChannelCheckpoint(ctx context.Context, req *datapb.UpdateChannelCheckpointRequest) (*commonpb.Status, error) { - return m.updateChanPos, m.err -} - -func (m *MockDataCoord) SaveImportSegment(ctx context.Context, req *datapb.SaveImportSegmentRequest) (*commonpb.Status, error) { - return m.addSegmentResp, m.err -} - -func (m *MockDataCoord) UnsetIsImportingState(context.Context, *datapb.UnsetIsImportingStateRequest) (*commonpb.Status, error) { - return m.unsetIsImportingStateResp, m.err -} - -func (m *MockDataCoord) MarkSegmentsDropped(ctx context.Context, req *datapb.MarkSegmentsDroppedRequest) (*commonpb.Status, error) { - return m.markSegmentsDroppedResp, m.err -} - -func (m *MockDataCoord) BroadcastAlteredCollection(ctx context.Context, req *datapb.AlterCollectionRequest) (*commonpb.Status, error) { - return m.broadCastResp, m.err -} - -func (m *MockDataCoord) CheckHealth(ctx context.Context, req *milvuspb.CheckHealthRequest) (*milvuspb.CheckHealthResponse, error) { - return &milvuspb.CheckHealthResponse{ - IsHealthy: true, - }, nil -} - -func (m *MockDataCoord) CreateIndex(ctx context.Context, req *indexpb.CreateIndexRequest) (*commonpb.Status, error) { - return m.createIndexResp, m.err -} - -func (m *MockDataCoord) DescribeIndex(ctx context.Context, req *indexpb.DescribeIndexRequest) (*indexpb.DescribeIndexResponse, error) { - return m.describeIndexResp, m.err -} - -func (m *MockDataCoord) GetIndexStatistics(ctx context.Context, req *indexpb.GetIndexStatisticsRequest) (*indexpb.GetIndexStatisticsResponse, error) { - return m.getIndexStatisticsResp, m.err -} - -func (m *MockDataCoord) GetIndexInfos(ctx context.Context, req *indexpb.GetIndexInfoRequest) (*indexpb.GetIndexInfoResponse, error) { - return m.getIndexInfosResp, m.err -} - -func (m *MockDataCoord) GetIndexState(ctx context.Context, req *indexpb.GetIndexStateRequest) (*indexpb.GetIndexStateResponse, error) { - return m.getIndexStateResp, m.err -} - -func (m *MockDataCoord) GetIndexBuildProgress(ctx context.Context, req *indexpb.GetIndexBuildProgressRequest) (*indexpb.GetIndexBuildProgressResponse, error) { - return m.getIndexBuildProgressResp, m.err -} - -func (m *MockDataCoord) GetSegmentIndexState(ctx context.Context, req *indexpb.GetSegmentIndexStateRequest) (*indexpb.GetSegmentIndexStateResponse, error) { - return m.getSegmentIndexStateResp, m.err -} - -func (m *MockDataCoord) DropIndex(ctx context.Context, req *indexpb.DropIndexRequest) (*commonpb.Status, error) { - return m.dropIndexResp, m.err -} - func Test_NewServer(t *testing.T) { paramtable.Init() - parameters := []string{"tikv", "etcd"} - for _, v := range parameters { - paramtable.Get().Save(paramtable.Get().MetaStoreCfg.MetaStoreType.Key, v) - ctx := context.Background() - getTiKVClient = func(cfg *paramtable.TiKVConfig) (*txnkv.Client, error) { - return tikv.SetupLocalTxn(), nil - } - defer func() { - getTiKVClient = tikv.GetTiKVClient - }() - server := NewServer(ctx, nil) - assert.NotNil(t, server) - t.Run("Run", func(t *testing.T) { - server.dataCoord = &MockDataCoord{} - // indexCoord := mocks.NewMockIndexCoord(t) - // indexCoord.EXPECT().Init().Return(nil) - // server.indexCoord = indexCoord + ctx := context.Background() + mockDataCoord := mocks.NewMockDataCoord(t) + server := NewServer(ctx, nil) + assert.NotNil(t, server) + server.dataCoord = mockDataCoord - err := server.Run() - assert.NoError(t, err) - }) - - t.Run("GetComponentStates", func(t *testing.T) { - server.dataCoord = &MockDataCoord{ - states: &milvuspb.ComponentStates{}, - } - states, err := server.GetComponentStates(ctx, nil) - assert.NoError(t, err) - assert.NotNil(t, states) - }) - - t.Run("GetTimeTickChannel", func(t *testing.T) { - server.dataCoord = &MockDataCoord{ - strResp: &milvuspb.StringResponse{}, - } - resp, err := server.GetTimeTickChannel(ctx, nil) - assert.NoError(t, err) - assert.NotNil(t, resp) - }) - - t.Run("GetStatisticsChannel", func(t *testing.T) { - server.dataCoord = &MockDataCoord{ - strResp: &milvuspb.StringResponse{}, - } - resp, err := server.GetStatisticsChannel(ctx, nil) - assert.NoError(t, err) - assert.NotNil(t, resp) - }) - - t.Run("GetSegmentInfo", func(t *testing.T) { - server.dataCoord = &MockDataCoord{ - infoResp: &datapb.GetSegmentInfoResponse{}, - } - resp, err := server.GetSegmentInfo(ctx, nil) - assert.NoError(t, err) - assert.NotNil(t, resp) - }) - - t.Run("Flush", func(t *testing.T) { - server.dataCoord = &MockDataCoord{ - flushResp: &datapb.FlushResponse{}, - } - resp, err := server.Flush(ctx, nil) - assert.NoError(t, err) - assert.NotNil(t, resp) - }) - - t.Run("AssignSegmentID", func(t *testing.T) { - server.dataCoord = &MockDataCoord{ - assignResp: &datapb.AssignSegmentIDResponse{}, - } - resp, err := server.AssignSegmentID(ctx, nil) - assert.NoError(t, err) - assert.NotNil(t, resp) - }) - - t.Run("GetSegmentStates", func(t *testing.T) { - server.dataCoord = &MockDataCoord{ - segStateResp: &datapb.GetSegmentStatesResponse{}, - } - resp, err := server.GetSegmentStates(ctx, nil) - assert.NoError(t, err) - assert.NotNil(t, resp) - }) - - t.Run("GetInsertBinlogPaths", func(t *testing.T) { - server.dataCoord = &MockDataCoord{ - binResp: &datapb.GetInsertBinlogPathsResponse{}, - } - resp, err := server.GetInsertBinlogPaths(ctx, nil) - assert.NoError(t, err) - assert.NotNil(t, resp) - }) - - t.Run("GetCollectionStatistics", func(t *testing.T) { - server.dataCoord = &MockDataCoord{ - colStatResp: &datapb.GetCollectionStatisticsResponse{}, - } - resp, err := server.GetCollectionStatistics(ctx, nil) - assert.NoError(t, err) - assert.NotNil(t, resp) - }) - - t.Run("GetPartitionStatistics", func(t *testing.T) { - server.dataCoord = &MockDataCoord{ - partStatResp: &datapb.GetPartitionStatisticsResponse{}, - } - resp, err := server.GetPartitionStatistics(ctx, nil) - assert.NoError(t, err) - assert.NotNil(t, resp) - }) - - t.Run("GetSegmentInfoChannel", func(t *testing.T) { - server.dataCoord = &MockDataCoord{ - strResp: &milvuspb.StringResponse{}, - } - resp, err := server.GetSegmentInfoChannel(ctx, nil) - assert.NoError(t, err) - assert.NotNil(t, resp) - }) - - t.Run("SaveBinlogPaths", func(t *testing.T) { - server.dataCoord = &MockDataCoord{ - status: &commonpb.Status{}, - } - resp, err := server.SaveBinlogPaths(ctx, nil) - assert.NoError(t, err) - assert.NotNil(t, resp) - }) - - t.Run("GetRecoveryInfo", func(t *testing.T) { - server.dataCoord = &MockDataCoord{ - recoverResp: &datapb.GetRecoveryInfoResponse{}, - } - resp, err := server.GetRecoveryInfo(ctx, nil) - assert.NoError(t, err) - assert.NotNil(t, resp) - }) - - t.Run("GetFlushedSegments", func(t *testing.T) { - server.dataCoord = &MockDataCoord{ - flushSegResp: &datapb.GetFlushedSegmentsResponse{}, - } - resp, err := server.GetFlushedSegments(ctx, nil) - assert.NoError(t, err) - assert.NotNil(t, resp) - }) - - t.Run("ShowConfigurations", func(t *testing.T) { - server.dataCoord = &MockDataCoord{ - configResp: &internalpb.ShowConfigurationsResponse{}, - } - resp, err := server.ShowConfigurations(ctx, nil) - assert.NoError(t, err) - assert.NotNil(t, resp) - }) - - t.Run("GetMetrics", func(t *testing.T) { - server.dataCoord = &MockDataCoord{ - metricResp: &milvuspb.GetMetricsResponse{}, - } - resp, err := server.GetMetrics(ctx, nil) - assert.NoError(t, err) - assert.NotNil(t, resp) - }) - - t.Run("WatchChannels", func(t *testing.T) { - server.dataCoord = &MockDataCoord{ - watchChannelsResp: &datapb.WatchChannelsResponse{}, - } - resp, err := server.WatchChannels(ctx, nil) - assert.NoError(t, err) - assert.NotNil(t, resp) - }) - - t.Run("GetFlushState", func(t *testing.T) { - server.dataCoord = &MockDataCoord{ - getFlushStateResp: &milvuspb.GetFlushStateResponse{}, - } - resp, err := server.GetFlushState(ctx, nil) - assert.NoError(t, err) - assert.NotNil(t, resp) - }) - - t.Run("GetFlushAllState", func(t *testing.T) { - server.dataCoord = &MockDataCoord{ - getFlushAllStateResp: &milvuspb.GetFlushAllStateResponse{}, - } - resp, err := server.GetFlushAllState(ctx, nil) - assert.NoError(t, err) - assert.NotNil(t, resp) - }) - - t.Run("DropVirtualChannel", func(t *testing.T) { - server.dataCoord = &MockDataCoord{ - dropVChanResp: &datapb.DropVirtualChannelResponse{}, - } - resp, err := server.DropVirtualChannel(ctx, nil) - assert.NoError(t, err) - assert.NotNil(t, resp) - }) - - t.Run("ManualCompaction", func(t *testing.T) { - server.dataCoord = &MockDataCoord{ - manualCompactionResp: &milvuspb.ManualCompactionResponse{}, - } - resp, err := server.ManualCompaction(ctx, nil) - assert.NoError(t, err) - assert.NotNil(t, resp) - }) - - t.Run("GetCompactionState", func(t *testing.T) { - server.dataCoord = &MockDataCoord{ - compactionStateResp: &milvuspb.GetCompactionStateResponse{}, - } - resp, err := server.GetCompactionState(ctx, nil) - assert.NoError(t, err) - assert.NotNil(t, resp) - }) - - t.Run("GetCompactionStateWithPlans", func(t *testing.T) { - server.dataCoord = &MockDataCoord{ - compactionPlansResp: &milvuspb.GetCompactionPlansResponse{}, - } - resp, err := server.GetCompactionStateWithPlans(ctx, nil) - assert.NoError(t, err) - assert.NotNil(t, resp) - }) - - t.Run("set segment state", func(t *testing.T) { - server.dataCoord = &MockDataCoord{ - setSegmentStateResp: &datapb.SetSegmentStateResponse{}, - } - resp, err := server.SetSegmentState(ctx, nil) - assert.NoError(t, err) - assert.NotNil(t, resp) - }) - - t.Run("import", func(t *testing.T) { - server.dataCoord = &MockDataCoord{ - importResp: &datapb.ImportTaskResponse{ - Status: &commonpb.Status{}, - }, - } - resp, err := server.Import(ctx, nil) - assert.NoError(t, err) - assert.NotNil(t, resp) - }) - - t.Run("update seg stat", func(t *testing.T) { - server.dataCoord = &MockDataCoord{ - updateSegStatResp: merr.Success(), - } - resp, err := server.UpdateSegmentStatistics(ctx, nil) - assert.NoError(t, err) - assert.NotNil(t, resp) - }) - - t.Run("UpdateChannelCheckpoint", func(t *testing.T) { - server.dataCoord = &MockDataCoord{ - updateChanPos: merr.Success(), - } - resp, err := server.UpdateChannelCheckpoint(ctx, nil) - assert.NoError(t, err) - assert.NotNil(t, resp) - }) - - t.Run("save import segment", func(t *testing.T) { - server.dataCoord = &MockDataCoord{ - addSegmentResp: merr.Success(), - } - resp, err := server.SaveImportSegment(ctx, nil) - assert.NoError(t, err) - assert.NotNil(t, resp) - }) - - t.Run("unset isImporting state", func(t *testing.T) { - server.dataCoord = &MockDataCoord{ - unsetIsImportingStateResp: merr.Success(), - } - resp, err := server.UnsetIsImportingState(ctx, nil) - assert.NoError(t, err) - assert.NotNil(t, resp) - }) - - t.Run("mark segments dropped", func(t *testing.T) { - server.dataCoord = &MockDataCoord{ - markSegmentsDroppedResp: merr.Success(), - } - resp, err := server.MarkSegmentsDropped(ctx, nil) - assert.NoError(t, err) - assert.NotNil(t, resp) - }) - - t.Run("broadcast altered collection", func(t *testing.T) { - server.dataCoord = &MockDataCoord{ - broadCastResp: &commonpb.Status{}, - } - resp, err := server.BroadcastAlteredCollection(ctx, nil) - assert.NoError(t, err) - assert.NotNil(t, resp) - }) - - t.Run("CheckHealth", func(t *testing.T) { - server.dataCoord = &MockDataCoord{} - ret, err := server.CheckHealth(ctx, nil) - assert.NoError(t, err) - assert.Equal(t, true, ret.IsHealthy) - }) - - t.Run("CreateIndex", func(t *testing.T) { - server.dataCoord = &MockDataCoord{ - createIndexResp: &commonpb.Status{}, - } - ret, err := server.CreateIndex(ctx, nil) - assert.NoError(t, err) - assert.NotNil(t, ret) - }) - - t.Run("DescribeIndex", func(t *testing.T) { - server.dataCoord = &MockDataCoord{ - describeIndexResp: &indexpb.DescribeIndexResponse{}, - } - ret, err := server.DescribeIndex(ctx, nil) - assert.NoError(t, err) - assert.NotNil(t, ret) - }) - - t.Run("GetIndexStatistics", func(t *testing.T) { - server.dataCoord = &MockDataCoord{ - getIndexStatisticsResp: &indexpb.GetIndexStatisticsResponse{}, - } - ret, err := server.GetIndexStatistics(ctx, nil) - assert.NoError(t, err) - assert.NotNil(t, ret) - }) - - t.Run("DropIndex", func(t *testing.T) { - server.dataCoord = &MockDataCoord{ - dropIndexResp: &commonpb.Status{}, - } - ret, err := server.DropIndex(ctx, nil) - assert.NoError(t, err) - assert.NotNil(t, ret) - }) - - t.Run("GetIndexState", func(t *testing.T) { - server.dataCoord = &MockDataCoord{ - getIndexStateResp: &indexpb.GetIndexStateResponse{}, - } - ret, err := server.GetIndexState(ctx, nil) - assert.NoError(t, err) - assert.NotNil(t, ret) - }) - - t.Run("GetIndexBuildProgress", func(t *testing.T) { - server.dataCoord = &MockDataCoord{ - getIndexBuildProgressResp: &indexpb.GetIndexBuildProgressResponse{}, - } - ret, err := server.GetIndexBuildProgress(ctx, nil) - assert.NoError(t, err) - assert.NotNil(t, ret) - }) - - t.Run("GetSegmentIndexState", func(t *testing.T) { - server.dataCoord = &MockDataCoord{ - getSegmentIndexStateResp: &indexpb.GetSegmentIndexStateResponse{}, - } - ret, err := server.GetSegmentIndexState(ctx, nil) - assert.NoError(t, err) - assert.NotNil(t, ret) - }) - - t.Run("GetIndexInfos", func(t *testing.T) { - server.dataCoord = &MockDataCoord{ - getIndexInfosResp: &indexpb.GetIndexInfoResponse{}, - } - ret, err := server.GetIndexInfos(ctx, nil) - assert.NoError(t, err) - assert.NotNil(t, ret) - }) - - err := server.Stop() + t.Run("GetComponentStates", func(t *testing.T) { + mockDataCoord.EXPECT().GetComponentStates(mock.Anything, mock.Anything).Return(&milvuspb.ComponentStates{}, nil) + states, err := server.GetComponentStates(ctx, nil) assert.NoError(t, err) - } + assert.NotNil(t, states) + }) + + t.Run("GetTimeTickChannel", func(t *testing.T) { + mockDataCoord.EXPECT().GetTimeTickChannel(mock.Anything, mock.Anything).Return(&milvuspb.StringResponse{}, nil) + resp, err := server.GetTimeTickChannel(ctx, nil) + assert.NoError(t, err) + assert.NotNil(t, resp) + }) + + t.Run("GetStatisticsChannel", func(t *testing.T) { + mockDataCoord.EXPECT().GetStatisticsChannel(mock.Anything, mock.Anything).Return(&milvuspb.StringResponse{}, nil) + resp, err := server.GetStatisticsChannel(ctx, nil) + assert.NoError(t, err) + assert.NotNil(t, resp) + }) + + t.Run("GetSegmentInfo", func(t *testing.T) { + mockDataCoord.EXPECT().GetSegmentInfo(mock.Anything, mock.Anything).Return(&datapb.GetSegmentInfoResponse{}, nil) + resp, err := server.GetSegmentInfo(ctx, nil) + assert.NoError(t, err) + assert.NotNil(t, resp) + }) + + t.Run("Flush", func(t *testing.T) { + mockDataCoord.EXPECT().Flush(mock.Anything, mock.Anything).Return(&datapb.FlushResponse{}, nil) + resp, err := server.Flush(ctx, nil) + assert.NoError(t, err) + assert.NotNil(t, resp) + }) + + t.Run("AssignSegmentID", func(t *testing.T) { + mockDataCoord.EXPECT().AssignSegmentID(mock.Anything, mock.Anything).Return(&datapb.AssignSegmentIDResponse{}, nil) + resp, err := server.AssignSegmentID(ctx, nil) + assert.NoError(t, err) + assert.NotNil(t, resp) + }) + + t.Run("GetSegmentStates", func(t *testing.T) { + mockDataCoord.EXPECT().GetSegmentStates(mock.Anything, mock.Anything).Return(&datapb.GetSegmentStatesResponse{}, nil) + resp, err := server.GetSegmentStates(ctx, nil) + assert.NoError(t, err) + assert.NotNil(t, resp) + }) + + t.Run("GetInsertBinlogPaths", func(t *testing.T) { + mockDataCoord.EXPECT().GetInsertBinlogPaths(mock.Anything, mock.Anything).Return(&datapb.GetInsertBinlogPathsResponse{}, nil) + resp, err := server.GetInsertBinlogPaths(ctx, nil) + assert.NoError(t, err) + assert.NotNil(t, resp) + }) + + t.Run("GetCollectionStatistics", func(t *testing.T) { + mockDataCoord.EXPECT().GetCollectionStatistics(mock.Anything, mock.Anything).Return(&datapb.GetCollectionStatisticsResponse{}, nil) + resp, err := server.GetCollectionStatistics(ctx, nil) + assert.NoError(t, err) + assert.NotNil(t, resp) + }) + + t.Run("GetPartitionStatistics", func(t *testing.T) { + mockDataCoord.EXPECT().GetPartitionStatistics(mock.Anything, mock.Anything).Return(&datapb.GetPartitionStatisticsResponse{}, nil) + resp, err := server.GetPartitionStatistics(ctx, nil) + assert.NoError(t, err) + assert.NotNil(t, resp) + }) + + t.Run("GetSegmentInfoChannel", func(t *testing.T) { + mockDataCoord.EXPECT().GetSegmentInfoChannel(mock.Anything, mock.Anything).Return(&milvuspb.StringResponse{}, nil) + resp, err := server.GetSegmentInfoChannel(ctx, nil) + assert.NoError(t, err) + assert.NotNil(t, resp) + }) + + t.Run("SaveBinlogPaths", func(t *testing.T) { + mockDataCoord.EXPECT().SaveBinlogPaths(mock.Anything, mock.Anything).Return(merr.Success(), nil) + resp, err := server.SaveBinlogPaths(ctx, nil) + assert.NoError(t, err) + assert.NotNil(t, resp) + }) + + t.Run("GetRecoveryInfo", func(t *testing.T) { + mockDataCoord.EXPECT().GetRecoveryInfo(mock.Anything, mock.Anything).Return(&datapb.GetRecoveryInfoResponse{}, nil) + resp, err := server.GetRecoveryInfo(ctx, nil) + assert.NoError(t, err) + assert.NotNil(t, resp) + }) + + t.Run("GetFlushedSegments", func(t *testing.T) { + mockDataCoord.EXPECT().GetFlushedSegments(mock.Anything, mock.Anything).Return(&datapb.GetFlushedSegmentsResponse{}, nil) + resp, err := server.GetFlushedSegments(ctx, nil) + assert.NoError(t, err) + assert.NotNil(t, resp) + }) + + t.Run("ShowConfigurations", func(t *testing.T) { + mockDataCoord.EXPECT().ShowConfigurations(mock.Anything, mock.Anything).Return(&internalpb.ShowConfigurationsResponse{}, nil) + resp, err := server.ShowConfigurations(ctx, nil) + assert.NoError(t, err) + assert.NotNil(t, resp) + }) + + t.Run("GetMetrics", func(t *testing.T) { + mockDataCoord.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(&milvuspb.GetMetricsResponse{}, nil) + resp, err := server.GetMetrics(ctx, nil) + assert.NoError(t, err) + assert.NotNil(t, resp) + }) + + t.Run("WatchChannels", func(t *testing.T) { + mockDataCoord.EXPECT().WatchChannels(mock.Anything, mock.Anything).Return(&datapb.WatchChannelsResponse{}, nil) + resp, err := server.WatchChannels(ctx, nil) + assert.NoError(t, err) + assert.NotNil(t, resp) + }) + + t.Run("GetFlushState", func(t *testing.T) { + mockDataCoord.EXPECT().GetFlushState(mock.Anything, mock.Anything).Return(&milvuspb.GetFlushStateResponse{}, nil) + resp, err := server.GetFlushState(ctx, nil) + assert.NoError(t, err) + assert.NotNil(t, resp) + }) + + t.Run("GetFlushAllState", func(t *testing.T) { + mockDataCoord.EXPECT().GetFlushAllState(mock.Anything, mock.Anything).Return(&milvuspb.GetFlushAllStateResponse{}, nil) + resp, err := server.GetFlushAllState(ctx, nil) + assert.NoError(t, err) + assert.NotNil(t, resp) + }) + + t.Run("DropVirtualChannel", func(t *testing.T) { + mockDataCoord.EXPECT().DropVirtualChannel(mock.Anything, mock.Anything).Return(&datapb.DropVirtualChannelResponse{}, nil) + resp, err := server.DropVirtualChannel(ctx, nil) + assert.NoError(t, err) + assert.NotNil(t, resp) + }) + + t.Run("ManualCompaction", func(t *testing.T) { + mockDataCoord.EXPECT().ManualCompaction(mock.Anything, mock.Anything).Return(&milvuspb.ManualCompactionResponse{}, nil) + resp, err := server.ManualCompaction(ctx, nil) + assert.NoError(t, err) + assert.NotNil(t, resp) + }) + + t.Run("GetCompactionState", func(t *testing.T) { + mockDataCoord.EXPECT().GetCompactionState(mock.Anything, mock.Anything).Return(&milvuspb.GetCompactionStateResponse{}, nil) + resp, err := server.GetCompactionState(ctx, nil) + assert.NoError(t, err) + assert.NotNil(t, resp) + }) + + t.Run("GetCompactionStateWithPlans", func(t *testing.T) { + mockDataCoord.EXPECT().GetCompactionStateWithPlans(mock.Anything, mock.Anything).Return(&milvuspb.GetCompactionPlansResponse{}, nil) + resp, err := server.GetCompactionStateWithPlans(ctx, nil) + assert.NoError(t, err) + assert.NotNil(t, resp) + }) + + t.Run("SetSegmentState", func(t *testing.T) { + mockDataCoord.EXPECT().SetSegmentState(mock.Anything, mock.Anything).Return(&datapb.SetSegmentStateResponse{}, nil) + resp, err := server.SetSegmentState(ctx, nil) + assert.NoError(t, err) + assert.NotNil(t, resp) + }) + + t.Run("import", func(t *testing.T) { + mockDataCoord.EXPECT().Import(mock.Anything, mock.Anything).Return(&datapb.ImportTaskResponse{}, nil) + resp, err := server.Import(ctx, nil) + assert.NoError(t, err) + assert.NotNil(t, resp) + }) + + t.Run("UpdateSegmentStatistics", func(t *testing.T) { + mockDataCoord.EXPECT().UpdateSegmentStatistics(mock.Anything, mock.Anything).Return(merr.Success(), nil) + resp, err := server.UpdateSegmentStatistics(ctx, nil) + assert.NoError(t, err) + assert.NotNil(t, resp) + }) + + t.Run("UpdateChannelCheckpoint", func(t *testing.T) { + mockDataCoord.EXPECT().UpdateChannelCheckpoint(mock.Anything, mock.Anything).Return(merr.Success(), nil) + resp, err := server.UpdateChannelCheckpoint(ctx, nil) + assert.NoError(t, err) + assert.NotNil(t, resp) + }) + + t.Run("SaveImportSegment", func(t *testing.T) { + mockDataCoord.EXPECT().SaveImportSegment(mock.Anything, mock.Anything).Return(merr.Success(), nil) + resp, err := server.SaveImportSegment(ctx, nil) + assert.NoError(t, err) + assert.NotNil(t, resp) + }) + + t.Run("UnsetIsImportingState", func(t *testing.T) { + mockDataCoord.EXPECT().UnsetIsImportingState(mock.Anything, mock.Anything).Return(merr.Success(), nil) + resp, err := server.UnsetIsImportingState(ctx, nil) + assert.NoError(t, err) + assert.NotNil(t, resp) + }) + + t.Run("MarkSegmentsDropped", func(t *testing.T) { + mockDataCoord.EXPECT().MarkSegmentsDropped(mock.Anything, mock.Anything).Return(merr.Success(), nil) + resp, err := server.MarkSegmentsDropped(ctx, nil) + assert.NoError(t, err) + assert.NotNil(t, resp) + }) + + t.Run("BroadcastAlteredCollection", func(t *testing.T) { + mockDataCoord.EXPECT().BroadcastAlteredCollection(mock.Anything, mock.Anything).Return(merr.Success(), nil) + resp, err := server.BroadcastAlteredCollection(ctx, nil) + assert.NoError(t, err) + assert.NotNil(t, resp) + }) + + t.Run("CheckHealth", func(t *testing.T) { + mockDataCoord.EXPECT().CheckHealth(mock.Anything, mock.Anything).Return(&milvuspb.CheckHealthResponse{IsHealthy: true}, nil) + ret, err := server.CheckHealth(ctx, nil) + assert.NoError(t, err) + assert.Equal(t, true, ret.IsHealthy) + }) + + t.Run("CreateIndex", func(t *testing.T) { + mockDataCoord.EXPECT().CreateIndex(mock.Anything, mock.Anything).Return(merr.Success(), nil) + ret, err := server.CreateIndex(ctx, nil) + assert.NoError(t, err) + assert.NotNil(t, ret) + }) + + t.Run("DescribeIndex", func(t *testing.T) { + mockDataCoord.EXPECT().DescribeIndex(mock.Anything, mock.Anything).Return(&indexpb.DescribeIndexResponse{}, nil) + ret, err := server.DescribeIndex(ctx, nil) + assert.NoError(t, err) + assert.NotNil(t, ret) + }) + + t.Run("GetIndexStatistics", func(t *testing.T) { + mockDataCoord.EXPECT().GetIndexStatistics(mock.Anything, mock.Anything).Return(&indexpb.GetIndexStatisticsResponse{}, nil) + ret, err := server.GetIndexStatistics(ctx, nil) + assert.NoError(t, err) + assert.NotNil(t, ret) + }) + + t.Run("DropIndex", func(t *testing.T) { + mockDataCoord.EXPECT().DropIndex(mock.Anything, mock.Anything).Return(merr.Success(), nil) + ret, err := server.DropIndex(ctx, nil) + assert.NoError(t, err) + assert.NotNil(t, ret) + }) + + t.Run("GetIndexState", func(t *testing.T) { + mockDataCoord.EXPECT().GetIndexState(mock.Anything, mock.Anything).Return(&indexpb.GetIndexStateResponse{}, nil) + ret, err := server.GetIndexState(ctx, nil) + assert.NoError(t, err) + assert.NotNil(t, ret) + }) + + t.Run("GetIndexBuildProgress", func(t *testing.T) { + mockDataCoord.EXPECT().GetIndexBuildProgress(mock.Anything, mock.Anything).Return(&indexpb.GetIndexBuildProgressResponse{}, nil) + ret, err := server.GetIndexBuildProgress(ctx, nil) + assert.NoError(t, err) + assert.NotNil(t, ret) + }) + + t.Run("GetSegmentIndexState", func(t *testing.T) { + mockDataCoord.EXPECT().GetSegmentIndexState(mock.Anything, mock.Anything).Return(&indexpb.GetSegmentIndexStateResponse{}, nil) + ret, err := server.GetSegmentIndexState(ctx, nil) + assert.NoError(t, err) + assert.NotNil(t, ret) + }) + + t.Run("GetIndexInfos", func(t *testing.T) { + mockDataCoord.EXPECT().GetIndexInfos(mock.Anything, mock.Anything).Return(&indexpb.GetIndexInfoResponse{}, nil) + ret, err := server.GetIndexInfos(ctx, nil) + assert.NoError(t, err) + assert.NotNil(t, ret) + }) } func Test_Run(t *testing.T) { paramtable.Init() - parameters := []string{"tikv", "etcd"} - for _, v := range parameters { - t.Run(fmt.Sprintf("Run server with %s as metadata storage", v), func(t *testing.T) { + + t.Run("test run success", func(t *testing.T) { + parameters := []string{"tikv", "etcd"} + for _, v := range parameters { paramtable.Get().Save(paramtable.Get().MetaStoreCfg.MetaStoreType.Key, v) ctx := context.Background() getTiKVClient = func(cfg *paramtable.TiKVConfig) (*txnkv.Client, error) { @@ -696,33 +342,97 @@ func Test_Run(t *testing.T) { server := NewServer(ctx, nil) assert.NotNil(t, server) - server.dataCoord = &MockDataCoord{ - regErr: errors.New("error"), - } + mockDataCoord := mocks.NewMockDataCoord(t) + server.dataCoord = mockDataCoord + mockDataCoord.EXPECT().SetEtcdClient(mock.Anything) + mockDataCoord.EXPECT().SetAddress(mock.Anything) + mockDataCoord.EXPECT().SetTiKVClient(mock.Anything).Maybe() + mockDataCoord.EXPECT().Init().Return(nil) + mockDataCoord.EXPECT().Start().Return(nil) + mockDataCoord.EXPECT().Register().Return(nil) err := server.Run() - assert.Error(t, err) - - server.dataCoord = &MockDataCoord{ - startErr: errors.New("error"), - } - - err = server.Run() - assert.Error(t, err) - - server.dataCoord = &MockDataCoord{ - initErr: errors.New("error"), - } - - err = server.Run() - assert.Error(t, err) - - server.dataCoord = &MockDataCoord{ - stopErr: errors.New("error"), - } + assert.NoError(t, err) + mockDataCoord.EXPECT().Stop().Return(nil) err = server.Stop() - assert.Error(t, err) - }) - } + assert.NoError(t, err) + } + }) + + paramtable.Get().Save(paramtable.Get().MetaStoreCfg.MetaStoreType.Key, "etcd") + + t.Run("test init error", func(t *testing.T) { + ctx := context.Background() + server := NewServer(ctx, nil) + assert.NotNil(t, server) + mockDataCoord := mocks.NewMockDataCoord(t) + mockDataCoord.EXPECT().SetEtcdClient(mock.Anything) + mockDataCoord.EXPECT().SetAddress(mock.Anything) + mockDataCoord.EXPECT().Init().Return(errors.New("error")) + server.dataCoord = mockDataCoord + + err := server.Run() + assert.Error(t, err) + + mockDataCoord.EXPECT().Stop().Return(nil) + server.Stop() + }) + + t.Run("test register error", func(t *testing.T) { + ctx := context.Background() + server := NewServer(ctx, nil) + assert.NotNil(t, server) + mockDataCoord := mocks.NewMockDataCoord(t) + mockDataCoord.EXPECT().SetEtcdClient(mock.Anything) + mockDataCoord.EXPECT().SetAddress(mock.Anything) + mockDataCoord.EXPECT().Init().Return(nil) + mockDataCoord.EXPECT().Register().Return(errors.New("error")) + server.dataCoord = mockDataCoord + + err := server.Run() + assert.Error(t, err) + + mockDataCoord.EXPECT().Stop().Return(nil) + server.Stop() + }) + + t.Run("test start error", func(t *testing.T) { + ctx := context.Background() + server := NewServer(ctx, nil) + assert.NotNil(t, server) + mockDataCoord := mocks.NewMockDataCoord(t) + mockDataCoord.EXPECT().SetEtcdClient(mock.Anything) + mockDataCoord.EXPECT().SetAddress(mock.Anything) + mockDataCoord.EXPECT().Init().Return(nil) + mockDataCoord.EXPECT().Register().Return(nil) + mockDataCoord.EXPECT().Start().Return(errors.New("error")) + server.dataCoord = mockDataCoord + + err := server.Run() + assert.Error(t, err) + + mockDataCoord.EXPECT().Stop().Return(nil) + server.Stop() + }) + + t.Run("test stop error", func(t *testing.T) { + ctx := context.Background() + server := NewServer(ctx, nil) + assert.NotNil(t, server) + mockDataCoord := mocks.NewMockDataCoord(t) + mockDataCoord.EXPECT().SetEtcdClient(mock.Anything) + mockDataCoord.EXPECT().SetAddress(mock.Anything) + mockDataCoord.EXPECT().Init().Return(nil) + mockDataCoord.EXPECT().Register().Return(nil) + mockDataCoord.EXPECT().Start().Return(nil) + server.dataCoord = mockDataCoord + + err := server.Run() + assert.NoError(t, err) + + mockDataCoord.EXPECT().Stop().Return(errors.New("error")) + err = server.Stop() + assert.Error(t, err) + }) } diff --git a/internal/distributed/datanode/service_test.go b/internal/distributed/datanode/service_test.go index ba4d5de241..88110e4d19 100644 --- a/internal/distributed/datanode/service_test.go +++ b/internal/distributed/datanode/service_test.go @@ -18,16 +18,16 @@ package grpcdatanode import ( "context" - "fmt" "testing" "github.com/cockroachdb/errors" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" clientv3 "go.etcd.io/etcd/client/v3" - "google.golang.org/grpc" "github.com/milvus-io/milvus-proto/go-api/v2/commonpb" "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb" + "github.com/milvus-io/milvus/internal/mocks" "github.com/milvus-io/milvus/internal/proto/datapb" "github.com/milvus-io/milvus/internal/proto/internalpb" "github.com/milvus-io/milvus/internal/types" @@ -182,52 +182,6 @@ func (m *MockDataNode) DropImport(ctx context.Context, req *datapb.DropImportReq return m.status, m.err } -// ///////////////////////////////////////////////////////////////////////////////////////////////////////////////////// -type mockDataCoord struct { - types.DataCoordClient -} - -func (m *mockDataCoord) GetComponentStates(ctx context.Context, req *milvuspb.GetComponentStatesRequest, opts ...grpc.CallOption) (*milvuspb.ComponentStates, error) { - return &milvuspb.ComponentStates{ - State: &milvuspb.ComponentInfo{ - StateCode: commonpb.StateCode_Healthy, - }, - Status: merr.Success(), - SubcomponentStates: []*milvuspb.ComponentInfo{ - { - StateCode: commonpb.StateCode_Healthy, - }, - }, - }, nil -} - -func (m *mockDataCoord) Stop() error { - return fmt.Errorf("stop error") -} - -// ///////////////////////////////////////////////////////////////////////////////////////////////////////////////////// -type mockRootCoord struct { - types.RootCoordClient -} - -func (m *mockRootCoord) GetComponentStates(ctx context.Context, req *milvuspb.GetComponentStatesRequest, opts ...grpc.CallOption) (*milvuspb.ComponentStates, error) { - return &milvuspb.ComponentStates{ - State: &milvuspb.ComponentInfo{ - StateCode: commonpb.StateCode_Healthy, - }, - Status: merr.Success(), - SubcomponentStates: []*milvuspb.ComponentInfo{ - { - StateCode: commonpb.StateCode_Healthy, - }, - }, - }, nil -} - -func (m *mockRootCoord) Stop() error { - return fmt.Errorf("stop error") -} - // ///////////////////////////////////////////////////////////////////////////////////////////////////////////////////// func Test_NewServer(t *testing.T) { paramtable.Init() @@ -236,12 +190,36 @@ func Test_NewServer(t *testing.T) { assert.NoError(t, err) assert.NotNil(t, server) + mockRootCoord := mocks.NewMockRootCoordClient(t) + mockRootCoord.EXPECT().GetComponentStates(mock.Anything, mock.Anything).Return(&milvuspb.ComponentStates{ + State: &milvuspb.ComponentInfo{ + StateCode: commonpb.StateCode_Healthy, + }, + Status: merr.Success(), + SubcomponentStates: []*milvuspb.ComponentInfo{ + { + StateCode: commonpb.StateCode_Healthy, + }, + }, + }, nil) server.newRootCoordClient = func() (types.RootCoordClient, error) { - return &mockRootCoord{}, nil + return mockRootCoord, nil } + mockDataCoord := mocks.NewMockDataCoordClient(t) + mockDataCoord.EXPECT().GetComponentStates(mock.Anything, mock.Anything).Return(&milvuspb.ComponentStates{ + State: &milvuspb.ComponentInfo{ + StateCode: commonpb.StateCode_Healthy, + }, + Status: merr.Success(), + SubcomponentStates: []*milvuspb.ComponentInfo{ + { + StateCode: commonpb.StateCode_Healthy, + }, + }, + }, nil) server.newDataCoordClient = func() (types.DataCoordClient, error) { - return &mockDataCoord{}, nil + return mockDataCoord, nil } t.Run("Run", func(t *testing.T) { @@ -371,18 +349,42 @@ func Test_Run(t *testing.T) { assert.NoError(t, err) assert.NotNil(t, server) + mockRootCoord := mocks.NewMockRootCoordClient(t) + mockRootCoord.EXPECT().GetComponentStates(mock.Anything, mock.Anything).Return(&milvuspb.ComponentStates{ + State: &milvuspb.ComponentInfo{ + StateCode: commonpb.StateCode_Healthy, + }, + Status: merr.Success(), + SubcomponentStates: []*milvuspb.ComponentInfo{ + { + StateCode: commonpb.StateCode_Healthy, + }, + }, + }, nil) + server.newRootCoordClient = func() (types.RootCoordClient, error) { + return mockRootCoord, nil + } + + mockDataCoord := mocks.NewMockDataCoordClient(t) + mockDataCoord.EXPECT().GetComponentStates(mock.Anything, mock.Anything).Return(&milvuspb.ComponentStates{ + State: &milvuspb.ComponentInfo{ + StateCode: commonpb.StateCode_Healthy, + }, + Status: merr.Success(), + SubcomponentStates: []*milvuspb.ComponentInfo{ + { + StateCode: commonpb.StateCode_Healthy, + }, + }, + }, nil) + server.newDataCoordClient = func() (types.DataCoordClient, error) { + return mockDataCoord, nil + } + server.datanode = &MockDataNode{ regErr: errors.New("error"), } - server.newRootCoordClient = func() (types.RootCoordClient, error) { - return &mockRootCoord{}, nil - } - - server.newDataCoordClient = func() (types.DataCoordClient, error) { - return &mockDataCoord{}, nil - } - err = server.Run() assert.Error(t, err) diff --git a/internal/distributed/proxy/client/client_test.go b/internal/distributed/proxy/client/client_test.go index 1043bf8f53..59dd2930d3 100644 --- a/internal/distributed/proxy/client/client_test.go +++ b/internal/distributed/proxy/client/client_test.go @@ -21,12 +21,14 @@ import ( "testing" "time" - "github.com/cockroachdb/errors" "github.com/stretchr/testify/assert" - "google.golang.org/grpc" + "github.com/stretchr/testify/mock" + "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb" + "github.com/milvus-io/milvus/internal/mocks" + "github.com/milvus-io/milvus/internal/proto/internalpb" "github.com/milvus-io/milvus/internal/proto/proxypb" - "github.com/milvus-io/milvus/internal/util/mock" + "github.com/milvus-io/milvus/pkg/util/merr" "github.com/milvus-io/milvus/pkg/util/paramtable" ) @@ -42,102 +44,392 @@ func Test_NewClient(t *testing.T) { assert.NoError(t, err) assert.NotNil(t, client) - checkFunc := func(retNotNil bool) { - retCheck := func(notNil bool, ret interface{}, err error) { - if notNil { - assert.NotNil(t, ret) - assert.NoError(t, err) - } else { - assert.Nil(t, ret) - assert.Error(t, err) - } - } - - r1, err := client.GetComponentStates(ctx, nil) - retCheck(retNotNil, r1, err) - - r2, err := client.GetStatisticsChannel(ctx, nil) - retCheck(retNotNil, r2, err) - - r3, err := client.InvalidateCollectionMetaCache(ctx, nil) - retCheck(retNotNil, r3, err) - - r7, err := client.InvalidateCredentialCache(ctx, nil) - retCheck(retNotNil, r7, err) - - r8, err := client.UpdateCredentialCache(ctx, nil) - retCheck(retNotNil, r8, err) - - { - r, err := client.RefreshPolicyInfoCache(ctx, nil) - retCheck(retNotNil, r, err) - } - } - - client.grpcClient = &mock.GRPCClientBase[proxypb.ProxyClient]{ - GetGrpcClientErr: errors.New("dummy"), - } - - newFunc1 := func(cc *grpc.ClientConn) proxypb.ProxyClient { - return &mock.GrpcProxyClient{Err: nil} - } - client.grpcClient.SetNewGrpcClientFunc(newFunc1) - - checkFunc(false) - - client.grpcClient = &mock.GRPCClientBase[proxypb.ProxyClient]{ - GetGrpcClientErr: nil, - } - - newFunc2 := func(cc *grpc.ClientConn) proxypb.ProxyClient { - return &mock.GrpcProxyClient{Err: errors.New("dummy")} - } - client.grpcClient.SetNewGrpcClientFunc(newFunc2) - checkFunc(false) - - client.grpcClient = &mock.GRPCClientBase[proxypb.ProxyClient]{ - GetGrpcClientErr: nil, - } - - newFunc3 := func(cc *grpc.ClientConn) proxypb.ProxyClient { - return &mock.GrpcProxyClient{Err: nil} - } - client.grpcClient.SetNewGrpcClientFunc(newFunc3) - - checkFunc(true) - - // timeout - timeout := time.Nanosecond - shortCtx, shortCancel := context.WithTimeout(ctx, timeout) - defer shortCancel() - time.Sleep(timeout) - - retCheck := func(ret interface{}, err error) { - assert.Nil(t, ret) - assert.Error(t, err) - } - - r1Timeout, err := client.GetComponentStates(shortCtx, nil) - retCheck(r1Timeout, err) - - r2Timeout, err := client.GetStatisticsChannel(shortCtx, nil) - retCheck(r2Timeout, err) - - r3Timeout, err := client.InvalidateCollectionMetaCache(shortCtx, nil) - retCheck(r3Timeout, err) - - r7Timeout, err := client.InvalidateCredentialCache(shortCtx, nil) - retCheck(r7Timeout, err) - - r8Timeout, err := client.UpdateCredentialCache(shortCtx, nil) - retCheck(r8Timeout, err) - - { - rTimeout, err := client.RefreshPolicyInfoCache(shortCtx, nil) - retCheck(rTimeout, err) - } - // cleanup err = client.Close() assert.NoError(t, err) } + +func Test_GetComponentStates(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx, "test", 1) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockProxyClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[proxypb.ProxyClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(proxypb.ProxyClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().GetComponentStates(mock.Anything, mock.Anything).Return(&milvuspb.ComponentStates{ + Status: merr.Success(), + }, nil) + _, err = client.GetComponentStates(ctx, &milvuspb.GetComponentStatesRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().GetComponentStates(mock.Anything, mock.Anything).Return(&milvuspb.ComponentStates{ + Status: merr.Status(merr.ErrServiceNotReady), + }, nil) + + _, err = client.GetComponentStates(ctx, &milvuspb.GetComponentStatesRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.GetComponentStates(ctx, &milvuspb.GetComponentStatesRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_GetStatisticsChannel(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx, "test", 1) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockProxyClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[proxypb.ProxyClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(proxypb.ProxyClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().GetStatisticsChannel(mock.Anything, mock.Anything).Return(&milvuspb.StringResponse{ + Status: merr.Success(), + }, nil) + _, err = client.GetStatisticsChannel(ctx, &internalpb.GetStatisticsChannelRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().GetStatisticsChannel(mock.Anything, mock.Anything).Return(&milvuspb.StringResponse{ + Status: merr.Status(merr.ErrServiceNotReady), + }, nil) + + _, err = client.GetStatisticsChannel(ctx, &internalpb.GetStatisticsChannelRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.GetStatisticsChannel(ctx, &internalpb.GetStatisticsChannelRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_InvalidateCollectionMetaCache(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx, "test", 1) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockProxyClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[proxypb.ProxyClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().GetNodeID().Return(1) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(proxypb.ProxyClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().InvalidateCollectionMetaCache(mock.Anything, mock.Anything).Return(merr.Success(), nil) + _, err = client.InvalidateCollectionMetaCache(ctx, &proxypb.InvalidateCollMetaCacheRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().InvalidateCollectionMetaCache(mock.Anything, mock.Anything).Return(merr.Status(merr.ErrServiceNotReady), nil) + + _, err = client.InvalidateCollectionMetaCache(ctx, &proxypb.InvalidateCollMetaCacheRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.InvalidateCollectionMetaCache(ctx, &proxypb.InvalidateCollMetaCacheRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_InvalidateCredentialCache(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx, "test", 1) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockProxyClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[proxypb.ProxyClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().GetNodeID().Return(1) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(proxypb.ProxyClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().InvalidateCredentialCache(mock.Anything, mock.Anything).Return(merr.Success(), nil) + _, err = client.InvalidateCredentialCache(ctx, &proxypb.InvalidateCredCacheRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().InvalidateCredentialCache(mock.Anything, mock.Anything).Return(merr.Status(merr.ErrServiceNotReady), nil) + + _, err = client.InvalidateCredentialCache(ctx, &proxypb.InvalidateCredCacheRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.InvalidateCredentialCache(ctx, &proxypb.InvalidateCredCacheRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_UpdateCredentialCache(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx, "test", 1) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockProxyClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[proxypb.ProxyClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().GetNodeID().Return(1) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(proxypb.ProxyClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().UpdateCredentialCache(mock.Anything, mock.Anything).Return(merr.Success(), nil) + _, err = client.UpdateCredentialCache(ctx, &proxypb.UpdateCredCacheRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().UpdateCredentialCache(mock.Anything, mock.Anything).Return(merr.Status(merr.ErrServiceNotReady), nil) + + _, err = client.UpdateCredentialCache(ctx, &proxypb.UpdateCredCacheRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.UpdateCredentialCache(ctx, &proxypb.UpdateCredCacheRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_RefreshPolicyInfoCache(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx, "test", 1) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockProxyClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[proxypb.ProxyClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().GetNodeID().Return(1) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(proxypb.ProxyClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().RefreshPolicyInfoCache(mock.Anything, mock.Anything).Return(merr.Success(), nil) + _, err = client.RefreshPolicyInfoCache(ctx, &proxypb.RefreshPolicyInfoCacheRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().RefreshPolicyInfoCache(mock.Anything, mock.Anything).Return(merr.Status(merr.ErrServiceNotReady), nil) + + _, err = client.RefreshPolicyInfoCache(ctx, &proxypb.RefreshPolicyInfoCacheRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.RefreshPolicyInfoCache(ctx, &proxypb.RefreshPolicyInfoCacheRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_GetProxyMetrics(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx, "test", 1) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockProxyClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[proxypb.ProxyClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().GetNodeID().Return(1) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(proxypb.ProxyClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().GetProxyMetrics(mock.Anything, mock.Anything).Return(&milvuspb.GetMetricsResponse{Status: merr.Success()}, nil) + _, err = client.GetProxyMetrics(ctx, &milvuspb.GetMetricsRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().GetProxyMetrics(mock.Anything, mock.Anything).Return(&milvuspb.GetMetricsResponse{Status: merr.Status(merr.ErrServiceNotReady)}, nil) + + _, err = client.GetProxyMetrics(ctx, &milvuspb.GetMetricsRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.GetProxyMetrics(ctx, &milvuspb.GetMetricsRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_SetRates(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx, "test", 1) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockProxyClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[proxypb.ProxyClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().GetNodeID().Return(1) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(proxypb.ProxyClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().SetRates(mock.Anything, mock.Anything).Return(merr.Success(), nil) + _, err = client.SetRates(ctx, &proxypb.SetRatesRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().SetRates(mock.Anything, mock.Anything).Return(merr.Status(merr.ErrServiceNotReady), nil) + + _, err = client.SetRates(ctx, &proxypb.SetRatesRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.SetRates(ctx, &proxypb.SetRatesRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_ListClientInfos(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx, "test", 1) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockProxyClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[proxypb.ProxyClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().GetNodeID().Return(1) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(proxypb.ProxyClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().ListClientInfos(mock.Anything, mock.Anything).Return(&proxypb.ListClientInfosResponse{Status: merr.Success()}, nil) + _, err = client.ListClientInfos(ctx, &proxypb.ListClientInfosRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().ListClientInfos(mock.Anything, mock.Anything).Return(&proxypb.ListClientInfosResponse{Status: merr.Status(merr.ErrServiceNotReady)}, nil) + + _, err = client.ListClientInfos(ctx, &proxypb.ListClientInfosRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.ListClientInfos(ctx, &proxypb.ListClientInfosRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func Test_GetDdChannel(t *testing.T) { + paramtable.Init() + + ctx := context.Background() + client, err := NewClient(ctx, "test", 1) + assert.NoError(t, err) + assert.NotNil(t, client) + defer client.Close() + + mockProxy := mocks.NewMockProxyClient(t) + mockGrpcClient := mocks.NewMockGrpcClient[proxypb.ProxyClient](t) + mockGrpcClient.EXPECT().Close().Return(nil) + mockGrpcClient.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(proxypb.ProxyClient) (interface{}, error)) (interface{}, error) { + return f(mockProxy) + }) + client.grpcClient = mockGrpcClient + + // test success + mockProxy.EXPECT().GetDdChannel(mock.Anything, mock.Anything).Return(&milvuspb.StringResponse{Status: merr.Success()}, nil) + _, err = client.GetDdChannel(ctx, &internalpb.GetDdChannelRequest{}) + assert.Nil(t, err) + + // test return error code + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().GetDdChannel(mock.Anything, mock.Anything).Return(&milvuspb.StringResponse{Status: merr.Status(merr.ErrServiceNotReady)}, nil) + + _, err = client.GetDdChannel(ctx, &internalpb.GetDdChannelRequest{}) + assert.Nil(t, err) + + // test ctx done + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) + _, err = client.GetDdChannel(ctx, &internalpb.GetDdChannelRequest{}) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} diff --git a/internal/distributed/proxy/service_test.go b/internal/distributed/proxy/service_test.go index 993340fa58..1aedb9d1d1 100644 --- a/internal/distributed/proxy/service_test.go +++ b/internal/distributed/proxy/service_test.go @@ -33,7 +33,6 @@ import ( "github.com/gin-gonic/gin" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/mock" - clientv3 "go.etcd.io/etcd/client/v3" "go.uber.org/zap" "golang.org/x/sync/errgroup" "google.golang.org/grpc" @@ -47,10 +46,7 @@ import ( grpcproxyclient "github.com/milvus-io/milvus/internal/distributed/proxy/client" "github.com/milvus-io/milvus/internal/distributed/proxy/httpserver" "github.com/milvus-io/milvus/internal/mocks" - "github.com/milvus-io/milvus/internal/proto/internalpb" - "github.com/milvus-io/milvus/internal/proto/proxypb" "github.com/milvus-io/milvus/internal/proxy" - "github.com/milvus-io/milvus/internal/types" milvusmock "github.com/milvus-io/milvus/internal/util/mock" "github.com/milvus-io/milvus/pkg/log" "github.com/milvus-io/milvus/pkg/util/funcutil" @@ -65,457 +61,6 @@ func TestMain(m *testing.M) { os.Exit(code) } -// ///////////////////////////////////////////////////////////////////////////////////////////////////////////////////// -type MockBase struct { - mock.Mock - isMockGetComponentStatesOn bool -} - -func (m *MockBase) On(methodName string, arguments ...interface{}) *mock.Call { - if methodName == "GetComponentStates" { - m.isMockGetComponentStatesOn = true - } - return m.Mock.On(methodName, arguments...) -} - -func (m *MockBase) GetComponentStates(ctx context.Context, req *milvuspb.GetComponentStatesRequest) (*milvuspb.ComponentStates, error) { - if m.isMockGetComponentStatesOn { - ret1 := &milvuspb.ComponentStates{} - var ret2 error - args := m.Called(ctx) - arg1 := args.Get(0) - arg2 := args.Get(1) - if arg1 != nil { - ret1 = arg1.(*milvuspb.ComponentStates) - } - if arg2 != nil { - ret2 = arg2.(error) - } - return ret1, ret2 - } - return &milvuspb.ComponentStates{ - State: &milvuspb.ComponentInfo{StateCode: commonpb.StateCode_Healthy}, - Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}, - }, nil -} - -func (m *MockBase) GetTimeTickChannel(ctx context.Context, req *internalpb.GetTimeTickChannelRequest) (*milvuspb.StringResponse, error) { - return nil, nil -} - -func (m *MockBase) GetStatisticsChannel(ctx context.Context, req *internalpb.GetStatisticsChannelRequest) (*milvuspb.StringResponse, error) { - return nil, nil -} - -// ///////////////////////////////////////////////////////////////////////////////////////////////////////////////////// -type MockProxy struct { - MockBase - err error - initErr error - startErr error - stopErr error - regErr error - isMockOn bool -} - -func (m *MockProxy) DescribeAlias(ctx context.Context, request *milvuspb.DescribeAliasRequest) (*milvuspb.DescribeAliasResponse, error) { - return nil, nil -} - -func (m *MockProxy) ListAliases(ctx context.Context, request *milvuspb.ListAliasesRequest) (*milvuspb.ListAliasesResponse, error) { - return nil, nil -} - -func (m *MockProxy) GetVersion(ctx context.Context, request *milvuspb.GetVersionRequest) (*milvuspb.GetVersionResponse, error) { - return nil, nil -} - -func (m *MockProxy) ListIndexedSegment(ctx context.Context, request *federpb.ListIndexedSegmentRequest) (*federpb.ListIndexedSegmentResponse, error) { - return nil, nil -} - -func (m *MockProxy) DescribeSegmentIndexData(ctx context.Context, request *federpb.DescribeSegmentIndexDataRequest) (*federpb.DescribeSegmentIndexDataResponse, error) { - return nil, nil -} - -func (m *MockProxy) SetRootCoordClient(rootCoord types.RootCoordClient) { -} - -func (m *MockProxy) SetDataCoordClient(dataCoord types.DataCoordClient) { -} - -func (m *MockProxy) SetQueryCoordClient(queryCoord types.QueryCoordClient) { -} - -func (m *MockProxy) SetQueryNodeCreator(f func(ctx context.Context, addr string, nodeID int64) (types.QueryNodeClient, error)) { - panic("error") -} - -func (m *MockProxy) Init() error { - return m.initErr -} - -func (m *MockProxy) Start() error { - return m.startErr -} - -func (m *MockProxy) Stop() error { - return m.stopErr -} - -func (m *MockProxy) Register() error { - return m.regErr -} - -func (m *MockProxy) ListClientInfos(ctx context.Context, request *proxypb.ListClientInfosRequest) (*proxypb.ListClientInfosResponse, error) { - return nil, nil -} - -func (m *MockProxy) InvalidateCollectionMetaCache(ctx context.Context, request *proxypb.InvalidateCollMetaCacheRequest) (*commonpb.Status, error) { - return nil, nil -} - -func (m *MockProxy) CreateDatabase(ctx context.Context, in *milvuspb.CreateDatabaseRequest) (*commonpb.Status, error) { - return nil, nil -} - -func (m *MockProxy) DropDatabase(ctx context.Context, in *milvuspb.DropDatabaseRequest) (*commonpb.Status, error) { - return nil, nil -} - -func (m *MockProxy) ListDatabases(ctx context.Context, in *milvuspb.ListDatabasesRequest) (*milvuspb.ListDatabasesResponse, error) { - return nil, nil -} - -func (m *MockProxy) CreateCollection(ctx context.Context, request *milvuspb.CreateCollectionRequest) (*commonpb.Status, error) { - return nil, nil -} - -func (m *MockProxy) DropCollection(ctx context.Context, request *milvuspb.DropCollectionRequest) (*commonpb.Status, error) { - return nil, nil -} - -func (m *MockProxy) HasCollection(ctx context.Context, request *milvuspb.HasCollectionRequest) (*milvuspb.BoolResponse, error) { - return nil, nil -} - -func (m *MockProxy) LoadCollection(ctx context.Context, request *milvuspb.LoadCollectionRequest) (*commonpb.Status, error) { - return nil, nil -} - -func (m *MockProxy) ReleaseCollection(ctx context.Context, request *milvuspb.ReleaseCollectionRequest) (*commonpb.Status, error) { - return nil, nil -} - -func (m *MockProxy) DescribeCollection(ctx context.Context, request *milvuspb.DescribeCollectionRequest) (*milvuspb.DescribeCollectionResponse, error) { - return nil, nil -} - -func (m *MockProxy) GetCollectionStatistics(ctx context.Context, request *milvuspb.GetCollectionStatisticsRequest) (*milvuspb.GetCollectionStatisticsResponse, error) { - return nil, nil -} - -func (m *MockProxy) ShowCollections(ctx context.Context, request *milvuspb.ShowCollectionsRequest) (*milvuspb.ShowCollectionsResponse, error) { - return nil, nil -} - -func (m *MockProxy) AlterCollection(ctx context.Context, request *milvuspb.AlterCollectionRequest) (*commonpb.Status, error) { - return nil, nil -} - -func (m *MockProxy) CreatePartition(ctx context.Context, request *milvuspb.CreatePartitionRequest) (*commonpb.Status, error) { - return nil, nil -} - -func (m *MockProxy) DropPartition(ctx context.Context, request *milvuspb.DropPartitionRequest) (*commonpb.Status, error) { - return nil, nil -} - -func (m *MockProxy) HasPartition(ctx context.Context, request *milvuspb.HasPartitionRequest) (*milvuspb.BoolResponse, error) { - return nil, nil -} - -func (m *MockProxy) LoadPartitions(ctx context.Context, request *milvuspb.LoadPartitionsRequest) (*commonpb.Status, error) { - return nil, nil -} - -func (m *MockProxy) ReleasePartitions(ctx context.Context, request *milvuspb.ReleasePartitionsRequest) (*commonpb.Status, error) { - return nil, nil -} - -func (m *MockProxy) GetPartitionStatistics(ctx context.Context, request *milvuspb.GetPartitionStatisticsRequest) (*milvuspb.GetPartitionStatisticsResponse, error) { - return nil, nil -} - -func (m *MockProxy) ShowPartitions(ctx context.Context, request *milvuspb.ShowPartitionsRequest) (*milvuspb.ShowPartitionsResponse, error) { - return nil, nil -} - -func (m *MockProxy) GetLoadingProgress(ctx context.Context, request *milvuspb.GetLoadingProgressRequest) (*milvuspb.GetLoadingProgressResponse, error) { - return nil, nil -} - -func (m *MockProxy) GetLoadState(ctx context.Context, request *milvuspb.GetLoadStateRequest) (*milvuspb.GetLoadStateResponse, error) { - return nil, nil -} - -func (m *MockProxy) CreateIndex(ctx context.Context, request *milvuspb.CreateIndexRequest) (*commonpb.Status, error) { - return nil, nil -} - -func (m *MockProxy) DropIndex(ctx context.Context, request *milvuspb.DropIndexRequest) (*commonpb.Status, error) { - return nil, nil -} - -func (m *MockProxy) DescribeIndex(ctx context.Context, request *milvuspb.DescribeIndexRequest) (*milvuspb.DescribeIndexResponse, error) { - return nil, nil -} - -func (m *MockProxy) GetIndexStatistics(ctx context.Context, request *milvuspb.GetIndexStatisticsRequest) (*milvuspb.GetIndexStatisticsResponse, error) { - return nil, nil -} - -func (m *MockProxy) GetIndexBuildProgress(ctx context.Context, request *milvuspb.GetIndexBuildProgressRequest) (*milvuspb.GetIndexBuildProgressResponse, error) { - return nil, nil -} - -func (m *MockProxy) GetIndexState(ctx context.Context, request *milvuspb.GetIndexStateRequest) (*milvuspb.GetIndexStateResponse, error) { - return nil, nil -} - -func (m *MockProxy) Insert(ctx context.Context, request *milvuspb.InsertRequest) (*milvuspb.MutationResult, error) { - return nil, nil -} - -func (m *MockProxy) Delete(ctx context.Context, request *milvuspb.DeleteRequest) (*milvuspb.MutationResult, error) { - return nil, nil -} - -func (m *MockProxy) Upsert(ctx context.Context, request *milvuspb.UpsertRequest) (*milvuspb.MutationResult, error) { - return nil, nil -} - -func (m *MockProxy) Search(ctx context.Context, request *milvuspb.SearchRequest) (*milvuspb.SearchResults, error) { - return nil, nil -} - -func (m *MockProxy) Flush(ctx context.Context, request *milvuspb.FlushRequest) (*milvuspb.FlushResponse, error) { - return nil, nil -} - -func (m *MockProxy) Query(ctx context.Context, request *milvuspb.QueryRequest) (*milvuspb.QueryResults, error) { - return nil, nil -} - -func (m *MockProxy) CalcDistance(ctx context.Context, request *milvuspb.CalcDistanceRequest) (*milvuspb.CalcDistanceResults, error) { - return nil, nil -} - -func (m *MockProxy) FlushAll(ctx context.Context, request *milvuspb.FlushAllRequest) (*milvuspb.FlushAllResponse, error) { - return nil, nil -} - -func (m *MockProxy) GetDdChannel(ctx context.Context, request *internalpb.GetDdChannelRequest) (*milvuspb.StringResponse, error) { - return nil, nil -} - -func (m *MockProxy) GetPersistentSegmentInfo(ctx context.Context, request *milvuspb.GetPersistentSegmentInfoRequest) (*milvuspb.GetPersistentSegmentInfoResponse, error) { - return nil, nil -} - -func (m *MockProxy) GetQuerySegmentInfo(ctx context.Context, request *milvuspb.GetQuerySegmentInfoRequest) (*milvuspb.GetQuerySegmentInfoResponse, error) { - return nil, nil -} - -func (m *MockProxy) Dummy(ctx context.Context, request *milvuspb.DummyRequest) (*milvuspb.DummyResponse, error) { - return nil, nil -} - -func (m *MockProxy) RegisterLink(ctx context.Context, request *milvuspb.RegisterLinkRequest) (*milvuspb.RegisterLinkResponse, error) { - return nil, nil -} - -func (m *MockProxy) GetMetrics(ctx context.Context, request *milvuspb.GetMetricsRequest) (*milvuspb.GetMetricsResponse, error) { - return nil, nil -} - -func (m *MockProxy) LoadBalance(ctx context.Context, request *milvuspb.LoadBalanceRequest) (*commonpb.Status, error) { - return nil, nil -} - -func (m *MockProxy) CreateAlias(ctx context.Context, request *milvuspb.CreateAliasRequest) (*commonpb.Status, error) { - return nil, nil -} - -func (m *MockProxy) DropAlias(ctx context.Context, request *milvuspb.DropAliasRequest) (*commonpb.Status, error) { - return nil, nil -} - -func (m *MockProxy) AlterAlias(ctx context.Context, request *milvuspb.AlterAliasRequest) (*commonpb.Status, error) { - return nil, nil -} - -func (m *MockProxy) SetRates(ctx context.Context, request *proxypb.SetRatesRequest) (*commonpb.Status, error) { - return nil, nil -} - -func (m *MockProxy) GetProxyMetrics(ctx context.Context, request *milvuspb.GetMetricsRequest) (*milvuspb.GetMetricsResponse, error) { - return nil, nil -} - -func (m *MockProxy) GetRateLimiter() (types.Limiter, error) { - return nil, nil -} - -func (m *MockProxy) UpdateStateCode(stateCode commonpb.StateCode) { -} - -func (m *MockProxy) SetAddress(address string) { -} - -func (m *MockProxy) GetAddress() string { - return "" -} - -func (m *MockProxy) SetEtcdClient(etcdClient *clientv3.Client) { -} - -func (m *MockProxy) GetCompactionState(ctx context.Context, req *milvuspb.GetCompactionStateRequest) (*milvuspb.GetCompactionStateResponse, error) { - return nil, nil -} - -func (m *MockProxy) ManualCompaction(ctx context.Context, req *milvuspb.ManualCompactionRequest) (*milvuspb.ManualCompactionResponse, error) { - return nil, nil -} - -func (m *MockProxy) GetCompactionStateWithPlans(ctx context.Context, req *milvuspb.GetCompactionPlansRequest) (*milvuspb.GetCompactionPlansResponse, error) { - return nil, nil -} - -func (m *MockProxy) GetFlushState(ctx context.Context, req *milvuspb.GetFlushStateRequest) (*milvuspb.GetFlushStateResponse, error) { - return nil, nil -} - -func (m *MockProxy) GetFlushAllState(ctx context.Context, req *milvuspb.GetFlushAllStateRequest) (*milvuspb.GetFlushAllStateResponse, error) { - return nil, nil -} - -func (m *MockProxy) Import(ctx context.Context, req *milvuspb.ImportRequest) (*milvuspb.ImportResponse, error) { - return nil, nil -} - -func (m *MockProxy) GetImportState(ctx context.Context, req *milvuspb.GetImportStateRequest) (*milvuspb.GetImportStateResponse, error) { - return nil, nil -} - -func (m *MockProxy) ListImportTasks(ctx context.Context, in *milvuspb.ListImportTasksRequest) (*milvuspb.ListImportTasksResponse, error) { - return nil, nil -} - -func (m *MockProxy) GetReplicas(ctx context.Context, req *milvuspb.GetReplicasRequest) (*milvuspb.GetReplicasResponse, error) { - return nil, nil -} - -func (m *MockProxy) InvalidateCredentialCache(ctx context.Context, request *proxypb.InvalidateCredCacheRequest) (*commonpb.Status, error) { - return nil, nil -} - -func (m *MockProxy) UpdateCredentialCache(ctx context.Context, request *proxypb.UpdateCredCacheRequest) (*commonpb.Status, error) { - return nil, nil -} - -func (m *MockProxy) CreateCredential(ctx context.Context, req *milvuspb.CreateCredentialRequest) (*commonpb.Status, error) { - return nil, nil -} - -func (m *MockProxy) UpdateCredential(ctx context.Context, req *milvuspb.UpdateCredentialRequest) (*commonpb.Status, error) { - return nil, nil -} - -func (m *MockProxy) DeleteCredential(ctx context.Context, req *milvuspb.DeleteCredentialRequest) (*commonpb.Status, error) { - return nil, nil -} - -func (m *MockProxy) ListCredUsers(ctx context.Context, req *milvuspb.ListCredUsersRequest) (*milvuspb.ListCredUsersResponse, error) { - return nil, nil -} - -func (m *MockProxy) CreateRole(ctx context.Context, req *milvuspb.CreateRoleRequest) (*commonpb.Status, error) { - return nil, nil -} - -func (m *MockProxy) DropRole(ctx context.Context, req *milvuspb.DropRoleRequest) (*commonpb.Status, error) { - return nil, nil -} - -func (m *MockProxy) OperateUserRole(ctx context.Context, req *milvuspb.OperateUserRoleRequest) (*commonpb.Status, error) { - return nil, nil -} - -func (m *MockProxy) SelectRole(ctx context.Context, req *milvuspb.SelectRoleRequest) (*milvuspb.SelectRoleResponse, error) { - return nil, nil -} - -func (m *MockProxy) SelectUser(ctx context.Context, req *milvuspb.SelectUserRequest) (*milvuspb.SelectUserResponse, error) { - return nil, nil -} - -func (m *MockProxy) OperatePrivilege(ctx context.Context, req *milvuspb.OperatePrivilegeRequest) (*commonpb.Status, error) { - return nil, nil -} - -func (m *MockProxy) SelectGrant(ctx context.Context, in *milvuspb.SelectGrantRequest) (*milvuspb.SelectGrantResponse, error) { - return nil, nil -} - -func (m *MockProxy) RefreshPolicyInfoCache(ctx context.Context, req *proxypb.RefreshPolicyInfoCacheRequest) (*commonpb.Status, error) { - return nil, nil -} - -func (m *MockProxy) CheckHealth(ctx context.Context, request *milvuspb.CheckHealthRequest) (*milvuspb.CheckHealthResponse, error) { - return nil, nil -} - -func (m *MockProxy) CreateResourceGroup(ctx context.Context, req *milvuspb.CreateResourceGroupRequest) (*commonpb.Status, error) { - return nil, nil -} - -func (m *MockProxy) DropResourceGroup(ctx context.Context, req *milvuspb.DropResourceGroupRequest) (*commonpb.Status, error) { - return nil, nil -} - -func (m *MockProxy) DescribeResourceGroup(ctx context.Context, req *milvuspb.DescribeResourceGroupRequest) (*milvuspb.DescribeResourceGroupResponse, error) { - return nil, nil -} - -func (m *MockProxy) TransferNode(ctx context.Context, req *milvuspb.TransferNodeRequest) (*commonpb.Status, error) { - return nil, nil -} - -func (m *MockProxy) TransferReplica(ctx context.Context, req *milvuspb.TransferReplicaRequest) (*commonpb.Status, error) { - return nil, nil -} - -func (m *MockProxy) ListResourceGroups(ctx context.Context, req *milvuspb.ListResourceGroupsRequest) (*milvuspb.ListResourceGroupsResponse, error) { - return nil, nil -} - -func (m *MockProxy) RenameCollection(ctx context.Context, req *milvuspb.RenameCollectionRequest) (*commonpb.Status, error) { - return nil, nil -} - -func (m *MockProxy) Connect(ctx context.Context, req *milvuspb.ConnectRequest) (*milvuspb.ConnectResponse, error) { - return nil, nil -} - -func (m *MockProxy) AllocTimestamp(ctx context.Context, req *milvuspb.AllocTimestampRequest) (*milvuspb.AllocTimestampResponse, error) { - return nil, nil -} - -func (m *MockProxy) ReplicateMessage(ctx context.Context, req *milvuspb.ReplicateMessageRequest) (*milvuspb.ReplicateMessageResponse, error) { - return nil, nil -} - // ///////////////////////////////////////////////////////////////////////////////////////////////////////////////////// type WaitOption struct { @@ -646,9 +191,25 @@ func Test_NewServer(t *testing.T) { ctx := context.Background() server := getServer(t) - var err error + assert.NotNil(t, server) + mockProxy := server.proxy.(*mocks.MockProxy) + t.Run("Run", func(t *testing.T) { - err = runAndWaitForServerReady(server) + mockProxy.EXPECT().Init().Return(nil) + mockProxy.EXPECT().Start().Return(nil) + mockProxy.EXPECT().Register().Return(nil) + mockProxy.EXPECT().SetEtcdClient(mock.Anything).Return() + mockProxy.EXPECT().GetRateLimiter().Return(nil, nil) + mockProxy.EXPECT().SetDataCoordClient(mock.Anything).Return() + mockProxy.EXPECT().SetRootCoordClient(mock.Anything).Return() + mockProxy.EXPECT().SetQueryCoordClient(mock.Anything).Return() + mockProxy.EXPECT().UpdateStateCode(mock.Anything).Return() + mockProxy.EXPECT().SetAddress(mock.Anything).Return() + err := runAndWaitForServerReady(server) + assert.NoError(t, err) + + mockProxy.EXPECT().Stop().Return(nil) + err = server.Stop() assert.NoError(t, err) }) @@ -658,387 +219,474 @@ func Test_NewServer(t *testing.T) { }) t.Run("GetStatisticsChannel", func(t *testing.T) { + mockProxy.EXPECT().GetStatisticsChannel(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.GetStatisticsChannel(ctx, nil) assert.NoError(t, err) }) t.Run("InvalidateCollectionMetaCache", func(t *testing.T) { + mockProxy.EXPECT().InvalidateCollectionMetaCache(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.InvalidateCollectionMetaCache(ctx, nil) assert.NoError(t, err) }) t.Run("CreateCollection", func(t *testing.T) { + mockProxy.EXPECT().CreateCollection(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.CreateCollection(ctx, nil) assert.NoError(t, err) }) t.Run("DropCollection", func(t *testing.T) { + mockProxy.EXPECT().DropCollection(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.DropCollection(ctx, nil) assert.NoError(t, err) }) t.Run("HasCollection", func(t *testing.T) { + mockProxy.EXPECT().HasCollection(mock.Anything, mock.Anything).Return(nil, nil) + mockProxy.EXPECT().HasCollection(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.HasCollection(ctx, nil) assert.NoError(t, err) }) t.Run("LoadCollection", func(t *testing.T) { + mockProxy.EXPECT().LoadCollection(mock.Anything, mock.Anything).Return(nil, nil) + mockProxy.EXPECT().LoadCollection(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.LoadCollection(ctx, nil) assert.NoError(t, err) }) t.Run("ReleaseCollection", func(t *testing.T) { + mockProxy.EXPECT().ReleaseCollection(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.ReleaseCollection(ctx, nil) assert.NoError(t, err) }) t.Run("DescribeCollection", func(t *testing.T) { + mockProxy.EXPECT().DescribeCollection(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.DescribeCollection(ctx, nil) assert.NoError(t, err) }) t.Run("GetCollectionStatistics", func(t *testing.T) { + mockProxy.EXPECT().GetCollectionStatistics(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.GetCollectionStatistics(ctx, nil) assert.NoError(t, err) }) t.Run("ShowCollections", func(t *testing.T) { + mockProxy.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.ShowCollections(ctx, nil) assert.NoError(t, err) }) t.Run("CreatePartition", func(t *testing.T) { + mockProxy.EXPECT().CreatePartition(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.CreatePartition(ctx, nil) assert.NoError(t, err) }) t.Run("DropPartition", func(t *testing.T) { + mockProxy.EXPECT().DropPartition(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.DropPartition(ctx, nil) assert.NoError(t, err) }) t.Run("HasPartition", func(t *testing.T) { + mockProxy.EXPECT().HasPartition(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.HasPartition(ctx, nil) assert.NoError(t, err) }) t.Run("LoadPartitions", func(t *testing.T) { + mockProxy.EXPECT().LoadPartitions(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.LoadPartitions(ctx, nil) assert.NoError(t, err) }) t.Run("ReleasePartitions", func(t *testing.T) { + mockProxy.EXPECT().ReleasePartitions(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.ReleasePartitions(ctx, nil) assert.NoError(t, err) }) t.Run("GetPartitionStatistics", func(t *testing.T) { + mockProxy.EXPECT().GetPartitionStatistics(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.GetPartitionStatistics(ctx, nil) assert.NoError(t, err) }) t.Run("ShowPartitions", func(t *testing.T) { + mockProxy.EXPECT().ShowPartitions(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.ShowPartitions(ctx, nil) assert.NoError(t, err) }) t.Run("GetLoadingProgress", func(t *testing.T) { + mockProxy.EXPECT().GetLoadingProgress(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.GetLoadingProgress(ctx, nil) assert.NoError(t, err) }) t.Run("CreateIndex", func(t *testing.T) { + mockProxy.EXPECT().CreateIndex(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.CreateIndex(ctx, nil) assert.NoError(t, err) }) t.Run("DropIndex", func(t *testing.T) { + mockProxy.EXPECT().DropIndex(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.DropIndex(ctx, nil) assert.NoError(t, err) }) t.Run("DescribeIndex", func(t *testing.T) { + mockProxy.EXPECT().DescribeIndex(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.DescribeIndex(ctx, nil) assert.NoError(t, err) }) t.Run("GetIndexStatistics", func(t *testing.T) { + mockProxy.EXPECT().GetIndexStatistics(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.GetIndexStatistics(ctx, nil) assert.NoError(t, err) }) t.Run("GetIndexBuildProgress", func(t *testing.T) { + mockProxy.EXPECT().GetIndexBuildProgress(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.GetIndexBuildProgress(ctx, nil) assert.NoError(t, err) }) t.Run("GetIndexState", func(t *testing.T) { + mockProxy.EXPECT().GetIndexState(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.GetIndexState(ctx, nil) assert.NoError(t, err) }) t.Run("Insert", func(t *testing.T) { + mockProxy.EXPECT().Insert(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.Insert(ctx, nil) assert.NoError(t, err) }) t.Run("Delete", func(t *testing.T) { + mockProxy.EXPECT().Delete(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.Delete(ctx, nil) assert.NoError(t, err) }) t.Run("Upsert", func(t *testing.T) { + mockProxy.EXPECT().Upsert(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.Upsert(ctx, nil) assert.NoError(t, err) }) t.Run("Search", func(t *testing.T) { + mockProxy.EXPECT().Search(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.Search(ctx, nil) assert.NoError(t, err) }) t.Run("Flush", func(t *testing.T) { + mockProxy.EXPECT().Flush(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.Flush(ctx, nil) assert.NoError(t, err) }) t.Run("Query", func(t *testing.T) { + mockProxy.EXPECT().Query(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.Query(ctx, nil) assert.NoError(t, err) }) t.Run("CalcDistance", func(t *testing.T) { + mockProxy.EXPECT().CalcDistance(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.CalcDistance(ctx, nil) assert.NoError(t, err) }) t.Run("GetDdChannel", func(t *testing.T) { + mockProxy.EXPECT().GetDdChannel(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.GetDdChannel(ctx, nil) assert.NoError(t, err) }) t.Run("GetPersistentSegmentInfo", func(t *testing.T) { + mockProxy.EXPECT().GetPersistentSegmentInfo(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.GetPersistentSegmentInfo(ctx, nil) assert.NoError(t, err) }) t.Run("GetQuerySegmentInfo", func(t *testing.T) { + mockProxy.EXPECT().GetQuerySegmentInfo(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.GetQuerySegmentInfo(ctx, nil) assert.NoError(t, err) }) t.Run("Dummy", func(t *testing.T) { + mockProxy.EXPECT().Dummy(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.Dummy(ctx, nil) assert.NoError(t, err) }) t.Run("RegisterLink", func(t *testing.T) { + mockProxy.EXPECT().RegisterLink(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.RegisterLink(ctx, nil) assert.NoError(t, err) }) t.Run("GetMetrics", func(t *testing.T) { + mockProxy.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.GetMetrics(ctx, nil) assert.NoError(t, err) }) t.Run("LoadBalance", func(t *testing.T) { + mockProxy.EXPECT().LoadBalance(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.LoadBalance(ctx, nil) assert.NoError(t, err) }) t.Run("CreateAlias", func(t *testing.T) { + mockProxy.EXPECT().CreateAlias(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.CreateAlias(ctx, nil) assert.NoError(t, err) }) t.Run("DropAlias", func(t *testing.T) { + mockProxy.EXPECT().DropAlias(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.DropAlias(ctx, nil) assert.NoError(t, err) }) t.Run("AlterAlias", func(t *testing.T) { + mockProxy.EXPECT().AlterAlias(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.AlterAlias(ctx, nil) assert.NoError(t, err) }) t.Run("GetCompactionState", func(t *testing.T) { + mockProxy.EXPECT().GetCompactionState(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.GetCompactionState(ctx, nil) assert.NoError(t, err) }) t.Run("ManualCompaction", func(t *testing.T) { + mockProxy.EXPECT().ManualCompaction(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.ManualCompaction(ctx, nil) assert.NoError(t, err) }) t.Run("GetCompactionStateWithPlans", func(t *testing.T) { + mockProxy.EXPECT().GetCompactionStateWithPlans(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.GetCompactionStateWithPlans(ctx, nil) assert.NoError(t, err) }) t.Run("CreateCredential", func(t *testing.T) { + mockProxy.EXPECT().CreateCredential(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.CreateCredential(ctx, nil) assert.NoError(t, err) }) t.Run("UpdateCredential", func(t *testing.T) { + mockProxy.EXPECT().UpdateCredential(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.UpdateCredential(ctx, nil) assert.NoError(t, err) }) t.Run("DeleteCredential", func(t *testing.T) { + mockProxy.EXPECT().DeleteCredential(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.DeleteCredential(ctx, nil) assert.NoError(t, err) }) t.Run("ListCredUsers", func(t *testing.T) { + mockProxy.EXPECT().ListCredUsers(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.ListCredUsers(ctx, nil) assert.NoError(t, err) }) t.Run("InvalidateCredentialCache", func(t *testing.T) { + mockProxy.EXPECT().InvalidateCredentialCache(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.InvalidateCredentialCache(ctx, nil) assert.NoError(t, err) }) t.Run("UpdateCredentialCache", func(t *testing.T) { + mockProxy.EXPECT().UpdateCredentialCache(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.UpdateCredentialCache(ctx, nil) assert.NoError(t, err) }) t.Run("CreateRole", func(t *testing.T) { + mockProxy.EXPECT().CreateRole(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.CreateRole(ctx, nil) assert.NoError(t, err) }) t.Run("DropRole", func(t *testing.T) { + mockProxy.EXPECT().DropRole(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.DropRole(ctx, nil) assert.NoError(t, err) }) t.Run("OperateUserRole", func(t *testing.T) { + mockProxy.EXPECT().OperateUserRole(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.OperateUserRole(ctx, nil) assert.NoError(t, err) }) t.Run("SelectRole", func(t *testing.T) { + mockProxy.EXPECT().SelectRole(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.SelectRole(ctx, nil) assert.NoError(t, err) }) t.Run("SelectUser", func(t *testing.T) { + mockProxy.EXPECT().SelectUser(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.SelectUser(ctx, nil) assert.NoError(t, err) }) t.Run("OperatePrivilege", func(t *testing.T) { + mockProxy.EXPECT().OperatePrivilege(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.OperatePrivilege(ctx, nil) assert.NoError(t, err) }) t.Run("SelectGrant", func(t *testing.T) { + mockProxy.EXPECT().SelectGrant(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.SelectGrant(ctx, nil) assert.NoError(t, err) }) t.Run("RefreshPrivilegeInfoCache", func(t *testing.T) { + mockProxy.EXPECT().RefreshPolicyInfoCache(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.RefreshPolicyInfoCache(ctx, nil) assert.NoError(t, err) }) t.Run("CheckHealth", func(t *testing.T) { + mockProxy.EXPECT().CheckHealth(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.CheckHealth(ctx, nil) assert.NoError(t, err) }) t.Run("RenameCollection", func(t *testing.T) { + mockProxy.EXPECT().RenameCollection(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.RenameCollection(ctx, nil) assert.NoError(t, err) }) t.Run("CreateResourceGroup", func(t *testing.T) { + mockProxy.EXPECT().CreateResourceGroup(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.CreateResourceGroup(ctx, nil) assert.NoError(t, err) }) t.Run("DropResourceGroup", func(t *testing.T) { + mockProxy.EXPECT().DropResourceGroup(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.DropResourceGroup(ctx, nil) assert.NoError(t, err) }) t.Run("TransferNode", func(t *testing.T) { + mockProxy.EXPECT().TransferNode(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.TransferNode(ctx, nil) assert.NoError(t, err) }) t.Run("TransferReplica", func(t *testing.T) { + mockProxy.EXPECT().TransferReplica(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.TransferReplica(ctx, nil) assert.NoError(t, err) }) t.Run("ListResourceGroups", func(t *testing.T) { + mockProxy.EXPECT().ListResourceGroups(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.ListResourceGroups(ctx, nil) assert.NoError(t, err) }) t.Run("DescribeResourceGroup", func(t *testing.T) { + mockProxy.EXPECT().DescribeResourceGroup(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.DescribeResourceGroup(ctx, nil) assert.NoError(t, err) }) t.Run("FlushAll", func(t *testing.T) { + mockProxy.EXPECT().FlushAll(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.FlushAll(ctx, nil) assert.NoError(t, err) }) t.Run("GetFlushAllState", func(t *testing.T) { + mockProxy.EXPECT().GetFlushAllState(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.GetFlushAllState(ctx, nil) assert.NoError(t, err) }) t.Run("CreateDatabase", func(t *testing.T) { + mockProxy.EXPECT().CreateDatabase(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.CreateDatabase(ctx, nil) assert.Nil(t, err) }) t.Run("DropDatabase", func(t *testing.T) { + mockProxy.EXPECT().DropDatabase(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.DropDatabase(ctx, nil) assert.Nil(t, err) }) t.Run("ListDatabase", func(t *testing.T) { + mockProxy.EXPECT().ListDatabases(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.ListDatabases(ctx, nil) assert.Nil(t, err) }) t.Run("AllocTimestamp", func(t *testing.T) { + mockProxy.EXPECT().AllocTimestamp(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.AllocTimestamp(ctx, nil) assert.Nil(t, err) }) - err = server.Stop() - assert.NoError(t, err) - // Update config and start server again to test with different config set. - // This works as config will be initialized only once - paramtable.Get().Save(proxy.Params.ProxyCfg.GinLogging.Key, "false") - err = runAndWaitForServerReady(server) - assert.NoError(t, err) - err = server.Stop() - assert.NoError(t, err) + t.Run("Run with different config", func(t *testing.T) { + mockProxy.EXPECT().Init().Return(nil) + mockProxy.EXPECT().Start().Return(nil) + mockProxy.EXPECT().Register().Return(nil) + mockProxy.EXPECT().SetEtcdClient(mock.Anything).Return() + mockProxy.EXPECT().GetRateLimiter().Return(nil, nil) + mockProxy.EXPECT().SetDataCoordClient(mock.Anything).Return() + mockProxy.EXPECT().SetRootCoordClient(mock.Anything).Return() + mockProxy.EXPECT().SetQueryCoordClient(mock.Anything).Return() + mockProxy.EXPECT().UpdateStateCode(mock.Anything).Return() + mockProxy.EXPECT().SetAddress(mock.Anything).Return() + // Update config and start server again to test with different config set. + // This works as config will be initialized only once + paramtable.Get().Save(proxy.Params.ProxyCfg.GinLogging.Key, "false") + err := runAndWaitForServerReady(server) + assert.NoError(t, err) + + mockProxy.EXPECT().Stop().Return(nil) + err = server.Stop() + assert.NoError(t, err) + }) } func TestServer_Check(t *testing.T) { ctx := context.Background() server := getServer(t) - mockProxy := server.proxy.(*MockProxy) + mockProxy := server.proxy.(*mocks.MockProxy) req := &grpc_health_v1.HealthCheckRequest{Service: ""} ret, err := server.Check(ctx, req) assert.NoError(t, err) assert.Equal(t, grpc_health_v1.HealthCheckResponse_SERVING, ret.Status) - mockProxy.On("GetComponentStates", ctx).Return(nil, fmt.Errorf("mock grpc unexpected error")).Once() + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().GetComponentStates(mock.Anything, mock.Anything).Return(nil, fmt.Errorf("mock grpc unexpected error")) ret, err = server.Check(ctx, req) assert.Error(t, err) @@ -1054,8 +702,9 @@ func TestServer_Check(t *testing.T) { State: componentInfo, Status: status, } - mockProxy.On("GetComponentStates", ctx).Return(componentState, nil) + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().GetComponentStates(mock.Anything, mock.Anything).Return(componentState, nil) ret, err = server.Check(ctx, req) assert.NoError(t, err) assert.Equal(t, grpc_health_v1.HealthCheckResponse_NOT_SERVING, ret.Status) @@ -1077,9 +726,8 @@ func TestServer_Check(t *testing.T) { } func TestServer_Watch(t *testing.T) { - ctx := context.Background() server := getServer(t) - mockProxy := server.proxy.(*MockProxy) + mockProxy := server.proxy.(*mocks.MockProxy) watchServer := milvusmock.NewGrpcHealthWatchServer() resultChan := watchServer.Chan() @@ -1091,7 +739,8 @@ func TestServer_Watch(t *testing.T) { assert.NoError(t, err) assert.Equal(t, grpc_health_v1.HealthCheckResponse_SERVING, ret.Status) - mockProxy.On("GetComponentStates", ctx).Return(nil, fmt.Errorf("mock grpc unexpected error")).Once() + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().GetComponentStates(mock.Anything, mock.Anything).Return(nil, fmt.Errorf("mock grpc unexpected error")) err = server.Watch(req, watchServer) ret = <-resultChan @@ -1108,7 +757,8 @@ func TestServer_Watch(t *testing.T) { State: componentInfo, Status: status, } - mockProxy.On("GetComponentStates", ctx).Return(componentState, nil) + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().GetComponentStates(mock.Anything, mock.Anything).Return(componentState, nil) err = server.Watch(req, watchServer) ret = <-resultChan @@ -1137,6 +787,19 @@ func TestServer_Watch(t *testing.T) { func Test_NewServer_HTTPServer_Enabled(t *testing.T) { server := getServer(t) + mockProxy := server.proxy.(*mocks.MockProxy) + mockProxy.EXPECT().Stop().Return(nil) + mockProxy.EXPECT().Init().Return(nil) + mockProxy.EXPECT().Start().Return(nil) + mockProxy.EXPECT().Register().Return(nil) + mockProxy.EXPECT().SetEtcdClient(mock.Anything).Return() + mockProxy.EXPECT().GetRateLimiter().Return(nil, nil) + mockProxy.EXPECT().SetDataCoordClient(mock.Anything).Return() + mockProxy.EXPECT().SetRootCoordClient(mock.Anything).Return() + mockProxy.EXPECT().SetQueryCoordClient(mock.Anything).Return() + mockProxy.EXPECT().UpdateStateCode(mock.Anything).Return() + mockProxy.EXPECT().SetAddress(mock.Anything).Return() + paramtable.Get().Save(proxy.Params.HTTPCfg.Enabled.Key, "true") err := runAndWaitForServerReady(server) assert.NoError(t, err) @@ -1159,11 +822,46 @@ func getServer(t *testing.T) *Server { assert.NotNil(t, server) assert.NoError(t, err) - server.proxy = &MockProxy{} - server.rootCoordClient = &milvusmock.GrpcRootCoordClient{} - server.dataCoordClient = &milvusmock.GrpcDataCoordClient{} + mockProxy := mocks.NewMockProxy(t) + mockProxy.EXPECT().GetComponentStates(mock.Anything, mock.Anything).Return(&milvuspb.ComponentStates{ + State: &milvuspb.ComponentInfo{ + NodeID: int64(uniquegenerator.GetUniqueIntGeneratorIns().GetInt()), + Role: "MockProxy", + StateCode: commonpb.StateCode_Healthy, + ExtraInfo: nil, + }, + SubcomponentStates: nil, + Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}, + }, nil).Maybe() + server.proxy = mockProxy - mockQC := &mocks.MockQueryCoordClient{} + mockRC := mocks.NewMockRootCoordClient(t) + mockRC.EXPECT().GetComponentStates(mock.Anything, mock.Anything, mock.Anything).Return(&milvuspb.ComponentStates{ + State: &milvuspb.ComponentInfo{ + NodeID: int64(uniquegenerator.GetUniqueIntGeneratorIns().GetInt()), + Role: "MockRootCoord", + StateCode: commonpb.StateCode_Healthy, + ExtraInfo: nil, + }, + SubcomponentStates: nil, + Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}, + }, nil).Maybe() + server.rootCoordClient = mockRC + + mockDC := mocks.NewMockDataCoordClient(t) + mockDC.EXPECT().GetComponentStates(mock.Anything, mock.Anything, mock.Anything).Return(&milvuspb.ComponentStates{ + State: &milvuspb.ComponentInfo{ + NodeID: int64(uniquegenerator.GetUniqueIntGeneratorIns().GetInt()), + Role: "MockDataCoord", + StateCode: commonpb.StateCode_Healthy, + ExtraInfo: nil, + }, + SubcomponentStates: nil, + Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}, + }, nil).Maybe() + server.dataCoordClient = mockDC + + mockQC := mocks.NewMockQueryCoordClient(t) server.queryCoordClient = mockQC mockQC.EXPECT().GetComponentStates(mock.Anything, mock.Anything, mock.Anything).Return(&milvuspb.ComponentStates{ State: &milvuspb.ComponentInfo{ @@ -1174,7 +872,7 @@ func getServer(t *testing.T) *Server { }, SubcomponentStates: nil, Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}, - }, nil) + }, nil).Maybe() return server } @@ -1182,6 +880,19 @@ func Test_NewServer_TLS_TwoWay(t *testing.T) { server := getServer(t) Params := ¶mtable.Get().ProxyGrpcServerCfg + mockProxy := server.proxy.(*mocks.MockProxy) + mockProxy.EXPECT().Stop().Return(nil) + mockProxy.EXPECT().Init().Return(nil) + mockProxy.EXPECT().Start().Return(nil) + mockProxy.EXPECT().Register().Return(nil) + mockProxy.EXPECT().SetEtcdClient(mock.Anything).Return() + mockProxy.EXPECT().GetRateLimiter().Return(nil, nil) + mockProxy.EXPECT().SetDataCoordClient(mock.Anything).Return() + mockProxy.EXPECT().SetRootCoordClient(mock.Anything).Return() + mockProxy.EXPECT().SetQueryCoordClient(mock.Anything).Return() + mockProxy.EXPECT().UpdateStateCode(mock.Anything).Return() + mockProxy.EXPECT().SetAddress(mock.Anything).Return() + paramtable.Get().Save(Params.TLSMode.Key, "2") paramtable.Get().Save(Params.ServerPemPath.Key, "../../../configs/cert/server.pem") paramtable.Get().Save(Params.ServerKeyPath.Key, "../../../configs/cert/server.key") @@ -1199,6 +910,19 @@ func Test_NewServer_TLS_OneWay(t *testing.T) { server := getServer(t) Params := ¶mtable.Get().ProxyGrpcServerCfg + mockProxy := server.proxy.(*mocks.MockProxy) + mockProxy.EXPECT().Stop().Return(nil) + mockProxy.EXPECT().Init().Return(nil) + mockProxy.EXPECT().Start().Return(nil) + mockProxy.EXPECT().Register().Return(nil) + mockProxy.EXPECT().SetEtcdClient(mock.Anything).Return() + mockProxy.EXPECT().GetRateLimiter().Return(nil, nil) + mockProxy.EXPECT().SetDataCoordClient(mock.Anything).Return() + mockProxy.EXPECT().SetRootCoordClient(mock.Anything).Return() + mockProxy.EXPECT().SetQueryCoordClient(mock.Anything).Return() + mockProxy.EXPECT().UpdateStateCode(mock.Anything).Return() + mockProxy.EXPECT().SetAddress(mock.Anything).Return() + paramtable.Get().Save(Params.TLSMode.Key, "1") paramtable.Get().Save(Params.ServerPemPath.Key, "../../../configs/cert/server.pem") paramtable.Get().Save(Params.ServerKeyPath.Key, "../../../configs/cert/server.key") @@ -1215,6 +939,12 @@ func Test_NewServer_TLS_FileNotExisted(t *testing.T) { server := getServer(t) Params := ¶mtable.Get().ProxyGrpcServerCfg + mockProxy := server.proxy.(*mocks.MockProxy) + mockProxy.EXPECT().Stop().Return(nil) + mockProxy.EXPECT().SetEtcdClient(mock.Anything).Return() + mockProxy.EXPECT().GetRateLimiter().Return(nil, nil) + mockProxy.EXPECT().SetAddress(mock.Anything).Return() + paramtable.Get().Save(Params.TLSMode.Key, "1") paramtable.Get().Save(Params.ServerPemPath.Key, "../not/existed/server.pem") paramtable.Get().Save(Params.ServerKeyPath.Key, "../../../configs/cert/server.key") @@ -1245,6 +975,19 @@ func Test_NewServer_TLS_FileNotExisted(t *testing.T) { func Test_NewHTTPServer_TLS_TwoWay(t *testing.T) { server := getServer(t) + mockProxy := server.proxy.(*mocks.MockProxy) + mockProxy.EXPECT().Stop().Return(nil) + mockProxy.EXPECT().Init().Return(nil) + mockProxy.EXPECT().Start().Return(nil) + mockProxy.EXPECT().Register().Return(nil) + mockProxy.EXPECT().SetEtcdClient(mock.Anything).Return() + mockProxy.EXPECT().GetRateLimiter().Return(nil, nil) + mockProxy.EXPECT().SetDataCoordClient(mock.Anything).Return() + mockProxy.EXPECT().SetRootCoordClient(mock.Anything).Return() + mockProxy.EXPECT().SetQueryCoordClient(mock.Anything).Return() + mockProxy.EXPECT().UpdateStateCode(mock.Anything).Return() + mockProxy.EXPECT().SetAddress(mock.Anything).Return() + Params := ¶mtable.Get().ProxyGrpcServerCfg paramtable.Get().Save(Params.TLSMode.Key, "2") @@ -1269,6 +1012,19 @@ func Test_NewHTTPServer_TLS_TwoWay(t *testing.T) { func Test_NewHTTPServer_TLS_OneWay(t *testing.T) { server := getServer(t) + mockProxy := server.proxy.(*mocks.MockProxy) + mockProxy.EXPECT().Stop().Return(nil) + mockProxy.EXPECT().Init().Return(nil) + mockProxy.EXPECT().Start().Return(nil) + mockProxy.EXPECT().Register().Return(nil) + mockProxy.EXPECT().SetEtcdClient(mock.Anything).Return() + mockProxy.EXPECT().GetRateLimiter().Return(nil, nil) + mockProxy.EXPECT().SetDataCoordClient(mock.Anything).Return() + mockProxy.EXPECT().SetRootCoordClient(mock.Anything).Return() + mockProxy.EXPECT().SetQueryCoordClient(mock.Anything).Return() + mockProxy.EXPECT().UpdateStateCode(mock.Anything).Return() + mockProxy.EXPECT().SetAddress(mock.Anything).Return() + Params := ¶mtable.Get().ProxyGrpcServerCfg paramtable.Get().Save(Params.TLSMode.Key, "1") @@ -1292,6 +1048,10 @@ func Test_NewHTTPServer_TLS_OneWay(t *testing.T) { func Test_NewHTTPServer_TLS_FileNotExisted(t *testing.T) { server := getServer(t) + mockProxy := server.proxy.(*mocks.MockProxy) + mockProxy.EXPECT().Stop().Return(nil) + mockProxy.EXPECT().SetEtcdClient(mock.Anything).Return() + mockProxy.EXPECT().SetAddress(mock.Anything).Return() Params := ¶mtable.Get().ProxyGrpcServerCfg paramtable.Get().Save(Params.TLSMode.Key, "1") @@ -1386,27 +1146,31 @@ func TestHttpAuthenticate(t *testing.T) { } func Test_Service_GracefulStop(t *testing.T) { - mockedProxy := mocks.NewMockProxy(t) var count int32 - mockedProxy.EXPECT().GetComponentStates(mock.Anything, mock.Anything).Run(func(_a0 context.Context, _a1 *milvuspb.GetComponentStatesRequest) { + server := getServer(t) + assert.NotNil(t, server) + + mockProxy := server.proxy.(*mocks.MockProxy) + mockProxy.ExpectedCalls = nil + mockProxy.EXPECT().GetComponentStates(mock.Anything, mock.Anything).Run(func(_a0 context.Context, _a1 *milvuspb.GetComponentStatesRequest) { fmt.Println("rpc start") time.Sleep(10 * time.Second) atomic.AddInt32(&count, 1) fmt.Println("rpc done") }).Return(&milvuspb.ComponentStates{Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}}, nil) - mockedProxy.EXPECT().Init().Return(nil) - mockedProxy.EXPECT().Start().Return(nil) - mockedProxy.EXPECT().Stop().Return(nil) - mockedProxy.EXPECT().Register().Return(nil) - mockedProxy.EXPECT().SetEtcdClient(mock.Anything).Return() - mockedProxy.EXPECT().GetRateLimiter().Return(nil, nil) - mockedProxy.EXPECT().SetDataCoordClient(mock.Anything).Return() - mockedProxy.EXPECT().SetRootCoordClient(mock.Anything).Return() - mockedProxy.EXPECT().SetQueryCoordClient(mock.Anything).Return() - mockedProxy.EXPECT().UpdateStateCode(mock.Anything).Return() - mockedProxy.EXPECT().SetAddress(mock.Anything).Return() + mockProxy.EXPECT().Init().Return(nil) + mockProxy.EXPECT().Start().Return(nil) + mockProxy.EXPECT().Stop().Return(nil) + mockProxy.EXPECT().Register().Return(nil) + mockProxy.EXPECT().SetEtcdClient(mock.Anything).Return() + mockProxy.EXPECT().GetRateLimiter().Return(nil, nil) + mockProxy.EXPECT().SetDataCoordClient(mock.Anything).Return() + mockProxy.EXPECT().SetRootCoordClient(mock.Anything).Return() + mockProxy.EXPECT().SetQueryCoordClient(mock.Anything).Return() + mockProxy.EXPECT().UpdateStateCode(mock.Anything).Return() + mockProxy.EXPECT().SetAddress(mock.Anything).Return() Params := ¶mtable.Get().ProxyGrpcServerCfg @@ -1426,10 +1190,6 @@ func Test_Service_GracefulStop(t *testing.T) { enableRegisterProxyServer = false }() - server := getServer(t) - assert.NotNil(t, server) - server.proxy = mockedProxy - err := server.Run() assert.Nil(t, err) diff --git a/internal/distributed/querycoord/service_test.go b/internal/distributed/querycoord/service_test.go index 5ad2ea0805..7c4f26f0b1 100644 --- a/internal/distributed/querycoord/service_test.go +++ b/internal/distributed/querycoord/service_test.go @@ -25,55 +25,17 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/mock" "github.com/tikv/client-go/v2/txnkv" - "google.golang.org/grpc" "github.com/milvus-io/milvus-proto/go-api/v2/commonpb" "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb" "github.com/milvus-io/milvus/internal/mocks" "github.com/milvus-io/milvus/internal/proto/internalpb" "github.com/milvus-io/milvus/internal/proto/querypb" - "github.com/milvus-io/milvus/internal/types" "github.com/milvus-io/milvus/pkg/util/merr" "github.com/milvus-io/milvus/pkg/util/paramtable" "github.com/milvus-io/milvus/pkg/util/tikv" ) -// ///////////////////////////////////////////////////////////////////////////////////////////////////////////////////// -type MockRootCoord struct { - types.RootCoordClient - stopErr error - stateErr commonpb.ErrorCode -} - -func (m *MockRootCoord) Close() error { - return m.stopErr -} - -func (m *MockRootCoord) GetComponentStates(ctx context.Context, req *milvuspb.GetComponentStatesRequest, opt ...grpc.CallOption) (*milvuspb.ComponentStates, error) { - return &milvuspb.ComponentStates{ - State: &milvuspb.ComponentInfo{StateCode: commonpb.StateCode_Healthy}, - Status: &commonpb.Status{ErrorCode: m.stateErr}, - }, nil -} - -// ///////////////////////////////////////////////////////////////////////////////////////////////////////////////////// -type MockDataCoord struct { - types.DataCoordClient - stopErr error - stateErr commonpb.ErrorCode -} - -func (m *MockDataCoord) Close() error { - return m.stopErr -} - -func (m *MockDataCoord) GetComponentStates(ctx context.Context, req *milvuspb.GetComponentStatesRequest, opts ...grpc.CallOption) (*milvuspb.ComponentStates, error) { - return &milvuspb.ComponentStates{ - State: &milvuspb.ComponentInfo{StateCode: commonpb.StateCode_Healthy}, - Status: &commonpb.Status{ErrorCode: m.stateErr}, - }, nil -} - // ///////////////////////////////////////////////////////////////////////////////////////////////////////////////////// func TestMain(m *testing.M) { paramtable.Init() @@ -96,13 +58,17 @@ func Test_NewServer(t *testing.T) { assert.NoError(t, err) assert.NotNil(t, server) - mdc := &MockDataCoord{ - stateErr: commonpb.ErrorCode_Success, - } + mdc := mocks.NewMockDataCoordClient(t) + mdc.EXPECT().GetComponentStates(mock.Anything, mock.Anything).Return(&milvuspb.ComponentStates{ + State: &milvuspb.ComponentInfo{StateCode: commonpb.StateCode_Healthy}, + Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}, + }, nil) - mrc := &MockRootCoord{ - stateErr: commonpb.ErrorCode_Success, - } + mrc := mocks.NewMockRootCoordClient(t) + mrc.EXPECT().GetComponentStates(mock.Anything, mock.Anything).Return(&milvuspb.ComponentStates{ + State: &milvuspb.ComponentInfo{StateCode: commonpb.StateCode_Healthy}, + Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}, + }, nil) mqc := getQueryCoord() successStatus := merr.Success() diff --git a/internal/distributed/rootcoord/service_test.go b/internal/distributed/rootcoord/service_test.go index d7b8e53e46..4df96aeafe 100644 --- a/internal/distributed/rootcoord/service_test.go +++ b/internal/distributed/rootcoord/service_test.go @@ -24,20 +24,18 @@ import ( "testing" "time" - "github.com/cockroachdb/errors" "github.com/stretchr/testify/assert" "github.com/tikv/client-go/v2/txnkv" clientv3 "go.etcd.io/etcd/client/v3" - "google.golang.org/grpc" "github.com/milvus-io/milvus-proto/go-api/v2/commonpb" "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb" + "github.com/milvus-io/milvus/internal/mocks" "github.com/milvus-io/milvus/internal/rootcoord" "github.com/milvus-io/milvus/internal/types" kvfactory "github.com/milvus-io/milvus/internal/util/dependency/kv" "github.com/milvus-io/milvus/internal/util/sessionutil" "github.com/milvus-io/milvus/pkg/util/etcd" - "github.com/milvus-io/milvus/pkg/util/merr" "github.com/milvus-io/milvus/pkg/util/paramtable" "github.com/milvus-io/milvus/pkg/util/tikv" ) @@ -109,42 +107,6 @@ func (m *mockCore) Stop() error { return fmt.Errorf("stop error") } -type mockDataCoord struct { - types.DataCoordClient -} - -func (m *mockDataCoord) Close() error { - return nil -} - -func (m *mockDataCoord) GetComponentStates(ctx context.Context, req *milvuspb.GetComponentStatesRequest, opts ...grpc.CallOption) (*milvuspb.ComponentStates, error) { - return &milvuspb.ComponentStates{ - State: &milvuspb.ComponentInfo{ - StateCode: commonpb.StateCode_Healthy, - }, - Status: merr.Success(), - SubcomponentStates: []*milvuspb.ComponentInfo{ - { - StateCode: commonpb.StateCode_Healthy, - }, - }, - }, nil -} - -func (m *mockDataCoord) Stop() error { - return fmt.Errorf("stop error") -} - -type mockQueryCoord struct { - types.QueryCoordClient - initErr error - startErr error -} - -func (m *mockQueryCoord) Close() error { - return fmt.Errorf("stop error") -} - func TestRun(t *testing.T) { paramtable.Init() parameters := []string{"tikv", "etcd"} @@ -169,11 +131,16 @@ func TestRun(t *testing.T) { assert.Error(t, err) assert.EqualError(t, err, "listen tcp: address 1000000: invalid port") + mockDataCoord := mocks.NewMockDataCoordClient(t) + mockDataCoord.EXPECT().Close().Return(nil) svr.newDataCoordClient = func() types.DataCoordClient { - return &mockDataCoord{} + return mockDataCoord } + + mockQueryCoord := mocks.NewMockQueryCoordClient(t) + mockQueryCoord.EXPECT().Close().Return(nil) svr.newQueryCoordClient = func() types.QueryCoordClient { - return &mockQueryCoord{} + return mockQueryCoord } paramtable.Get().Save(rcServerConfig.Port.Key, fmt.Sprintf("%d", rand.Int()%100+10000)) @@ -251,8 +218,10 @@ func TestServerRun_DataCoordClientInitErr(t *testing.T) { assert.NoError(t, err) assert.NotNil(t, server) + mockDataCoord := mocks.NewMockDataCoordClient(t) + mockDataCoord.EXPECT().Close().Return(nil) server.newDataCoordClient = func() types.DataCoordClient { - return &mockDataCoord{} + return mockDataCoord } assert.Panics(t, func() { server.Run() }) @@ -277,8 +246,10 @@ func TestServerRun_DataCoordClientStartErr(t *testing.T) { assert.NoError(t, err) assert.NotNil(t, server) + mockDataCoord := mocks.NewMockDataCoordClient(t) + mockDataCoord.EXPECT().Close().Return(nil) server.newDataCoordClient = func() types.DataCoordClient { - return &mockDataCoord{} + return mockDataCoord } assert.Panics(t, func() { server.Run() }) @@ -303,9 +274,12 @@ func TestServerRun_QueryCoordClientInitErr(t *testing.T) { assert.NoError(t, err) assert.NotNil(t, server) + mockQueryCoord := mocks.NewMockQueryCoordClient(t) + mockQueryCoord.EXPECT().Close().Return(nil) server.newQueryCoordClient = func() types.QueryCoordClient { - return &mockQueryCoord{initErr: errors.New("mock querycoord init error")} + return mockQueryCoord } + assert.Panics(t, func() { server.Run() }) err = server.Stop() @@ -329,8 +303,10 @@ func TestServer_QueryCoordClientStartErr(t *testing.T) { assert.NoError(t, err) assert.NotNil(t, server) + mockQueryCoord := mocks.NewMockQueryCoordClient(t) + mockQueryCoord.EXPECT().Close().Return(nil) server.newQueryCoordClient = func() types.QueryCoordClient { - return &mockQueryCoord{startErr: errors.New("mock querycoord start error")} + return mockQueryCoord } assert.Panics(t, func() { server.Run() }) diff --git a/internal/mocks/mock_grpc_client.go b/internal/mocks/mock_grpc_client.go new file mode 100644 index 0000000000..e47fa2bf40 --- /dev/null +++ b/internal/mocks/mock_grpc_client.go @@ -0,0 +1,472 @@ +// Code generated by mockery v2.32.4. DO NOT EDIT. + +package mocks + +import ( + context "context" + + grpc "google.golang.org/grpc" + + grpcclient "github.com/milvus-io/milvus/internal/util/grpcclient" + + mock "github.com/stretchr/testify/mock" + + sessionutil "github.com/milvus-io/milvus/internal/util/sessionutil" +) + +// MockGrpcClient is an autogenerated mock type for the GrpcClient type +type MockGrpcClient[T grpcclient.GrpcComponent] struct { + mock.Mock +} + +type MockGrpcClient_Expecter[T grpcclient.GrpcComponent] struct { + mock *mock.Mock +} + +func (_m *MockGrpcClient[T]) EXPECT() *MockGrpcClient_Expecter[T] { + return &MockGrpcClient_Expecter[T]{mock: &_m.Mock} +} + +// Call provides a mock function with given fields: ctx, caller +func (_m *MockGrpcClient[T]) Call(ctx context.Context, caller func(T) (interface{}, error)) (interface{}, error) { + ret := _m.Called(ctx, caller) + + var r0 interface{} + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, func(T) (interface{}, error)) (interface{}, error)); ok { + return rf(ctx, caller) + } + if rf, ok := ret.Get(0).(func(context.Context, func(T) (interface{}, error)) interface{}); ok { + r0 = rf(ctx, caller) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(interface{}) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, func(T) (interface{}, error)) error); ok { + r1 = rf(ctx, caller) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// MockGrpcClient_Call_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Call' +type MockGrpcClient_Call_Call[T grpcclient.GrpcComponent] struct { + *mock.Call +} + +// Call is a helper method to define mock.On call +// - ctx context.Context +// - caller func(T)(interface{} , error) +func (_e *MockGrpcClient_Expecter[T]) Call(ctx interface{}, caller interface{}) *MockGrpcClient_Call_Call[T] { + return &MockGrpcClient_Call_Call[T]{Call: _e.mock.On("Call", ctx, caller)} +} + +func (_c *MockGrpcClient_Call_Call[T]) Run(run func(ctx context.Context, caller func(T) (interface{}, error))) *MockGrpcClient_Call_Call[T] { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].(func(T) (interface{}, error))) + }) + return _c +} + +func (_c *MockGrpcClient_Call_Call[T]) Return(_a0 interface{}, _a1 error) *MockGrpcClient_Call_Call[T] { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MockGrpcClient_Call_Call[T]) RunAndReturn(run func(context.Context, func(T) (interface{}, error)) (interface{}, error)) *MockGrpcClient_Call_Call[T] { + _c.Call.Return(run) + return _c +} + +// Close provides a mock function with given fields: +func (_m *MockGrpcClient[T]) Close() error { + ret := _m.Called() + + var r0 error + if rf, ok := ret.Get(0).(func() error); ok { + r0 = rf() + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// MockGrpcClient_Close_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Close' +type MockGrpcClient_Close_Call[T grpcclient.GrpcComponent] struct { + *mock.Call +} + +// Close is a helper method to define mock.On call +func (_e *MockGrpcClient_Expecter[T]) Close() *MockGrpcClient_Close_Call[T] { + return &MockGrpcClient_Close_Call[T]{Call: _e.mock.On("Close")} +} + +func (_c *MockGrpcClient_Close_Call[T]) Run(run func()) *MockGrpcClient_Close_Call[T] { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *MockGrpcClient_Close_Call[T]) Return(_a0 error) *MockGrpcClient_Close_Call[T] { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockGrpcClient_Close_Call[T]) RunAndReturn(run func() error) *MockGrpcClient_Close_Call[T] { + _c.Call.Return(run) + return _c +} + +// EnableEncryption provides a mock function with given fields: +func (_m *MockGrpcClient[T]) EnableEncryption() { + _m.Called() +} + +// MockGrpcClient_EnableEncryption_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'EnableEncryption' +type MockGrpcClient_EnableEncryption_Call[T grpcclient.GrpcComponent] struct { + *mock.Call +} + +// EnableEncryption is a helper method to define mock.On call +func (_e *MockGrpcClient_Expecter[T]) EnableEncryption() *MockGrpcClient_EnableEncryption_Call[T] { + return &MockGrpcClient_EnableEncryption_Call[T]{Call: _e.mock.On("EnableEncryption")} +} + +func (_c *MockGrpcClient_EnableEncryption_Call[T]) Run(run func()) *MockGrpcClient_EnableEncryption_Call[T] { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *MockGrpcClient_EnableEncryption_Call[T]) Return() *MockGrpcClient_EnableEncryption_Call[T] { + _c.Call.Return() + return _c +} + +func (_c *MockGrpcClient_EnableEncryption_Call[T]) RunAndReturn(run func()) *MockGrpcClient_EnableEncryption_Call[T] { + _c.Call.Return(run) + return _c +} + +// GetNodeID provides a mock function with given fields: +func (_m *MockGrpcClient[T]) GetNodeID() int64 { + ret := _m.Called() + + var r0 int64 + if rf, ok := ret.Get(0).(func() int64); ok { + r0 = rf() + } else { + r0 = ret.Get(0).(int64) + } + + return r0 +} + +// MockGrpcClient_GetNodeID_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetNodeID' +type MockGrpcClient_GetNodeID_Call[T grpcclient.GrpcComponent] struct { + *mock.Call +} + +// GetNodeID is a helper method to define mock.On call +func (_e *MockGrpcClient_Expecter[T]) GetNodeID() *MockGrpcClient_GetNodeID_Call[T] { + return &MockGrpcClient_GetNodeID_Call[T]{Call: _e.mock.On("GetNodeID")} +} + +func (_c *MockGrpcClient_GetNodeID_Call[T]) Run(run func()) *MockGrpcClient_GetNodeID_Call[T] { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *MockGrpcClient_GetNodeID_Call[T]) Return(_a0 int64) *MockGrpcClient_GetNodeID_Call[T] { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockGrpcClient_GetNodeID_Call[T]) RunAndReturn(run func() int64) *MockGrpcClient_GetNodeID_Call[T] { + _c.Call.Return(run) + return _c +} + +// GetRole provides a mock function with given fields: +func (_m *MockGrpcClient[T]) GetRole() string { + ret := _m.Called() + + var r0 string + if rf, ok := ret.Get(0).(func() string); ok { + r0 = rf() + } else { + r0 = ret.Get(0).(string) + } + + return r0 +} + +// MockGrpcClient_GetRole_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetRole' +type MockGrpcClient_GetRole_Call[T grpcclient.GrpcComponent] struct { + *mock.Call +} + +// GetRole is a helper method to define mock.On call +func (_e *MockGrpcClient_Expecter[T]) GetRole() *MockGrpcClient_GetRole_Call[T] { + return &MockGrpcClient_GetRole_Call[T]{Call: _e.mock.On("GetRole")} +} + +func (_c *MockGrpcClient_GetRole_Call[T]) Run(run func()) *MockGrpcClient_GetRole_Call[T] { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *MockGrpcClient_GetRole_Call[T]) Return(_a0 string) *MockGrpcClient_GetRole_Call[T] { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockGrpcClient_GetRole_Call[T]) RunAndReturn(run func() string) *MockGrpcClient_GetRole_Call[T] { + _c.Call.Return(run) + return _c +} + +// ReCall provides a mock function with given fields: ctx, caller +func (_m *MockGrpcClient[T]) ReCall(ctx context.Context, caller func(T) (interface{}, error)) (interface{}, error) { + ret := _m.Called(ctx, caller) + + var r0 interface{} + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, func(T) (interface{}, error)) (interface{}, error)); ok { + return rf(ctx, caller) + } + if rf, ok := ret.Get(0).(func(context.Context, func(T) (interface{}, error)) interface{}); ok { + r0 = rf(ctx, caller) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(interface{}) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, func(T) (interface{}, error)) error); ok { + r1 = rf(ctx, caller) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// MockGrpcClient_ReCall_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'ReCall' +type MockGrpcClient_ReCall_Call[T grpcclient.GrpcComponent] struct { + *mock.Call +} + +// ReCall is a helper method to define mock.On call +// - ctx context.Context +// - caller func(T)(interface{} , error) +func (_e *MockGrpcClient_Expecter[T]) ReCall(ctx interface{}, caller interface{}) *MockGrpcClient_ReCall_Call[T] { + return &MockGrpcClient_ReCall_Call[T]{Call: _e.mock.On("ReCall", ctx, caller)} +} + +func (_c *MockGrpcClient_ReCall_Call[T]) Run(run func(ctx context.Context, caller func(T) (interface{}, error))) *MockGrpcClient_ReCall_Call[T] { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].(func(T) (interface{}, error))) + }) + return _c +} + +func (_c *MockGrpcClient_ReCall_Call[T]) Return(_a0 interface{}, _a1 error) *MockGrpcClient_ReCall_Call[T] { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MockGrpcClient_ReCall_Call[T]) RunAndReturn(run func(context.Context, func(T) (interface{}, error)) (interface{}, error)) *MockGrpcClient_ReCall_Call[T] { + _c.Call.Return(run) + return _c +} + +// SetGetAddrFunc provides a mock function with given fields: _a0 +func (_m *MockGrpcClient[T]) SetGetAddrFunc(_a0 func() (string, error)) { + _m.Called(_a0) +} + +// MockGrpcClient_SetGetAddrFunc_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SetGetAddrFunc' +type MockGrpcClient_SetGetAddrFunc_Call[T grpcclient.GrpcComponent] struct { + *mock.Call +} + +// SetGetAddrFunc is a helper method to define mock.On call +// - _a0 func()(string , error) +func (_e *MockGrpcClient_Expecter[T]) SetGetAddrFunc(_a0 interface{}) *MockGrpcClient_SetGetAddrFunc_Call[T] { + return &MockGrpcClient_SetGetAddrFunc_Call[T]{Call: _e.mock.On("SetGetAddrFunc", _a0)} +} + +func (_c *MockGrpcClient_SetGetAddrFunc_Call[T]) Run(run func(_a0 func() (string, error))) *MockGrpcClient_SetGetAddrFunc_Call[T] { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(func() (string, error))) + }) + return _c +} + +func (_c *MockGrpcClient_SetGetAddrFunc_Call[T]) Return() *MockGrpcClient_SetGetAddrFunc_Call[T] { + _c.Call.Return() + return _c +} + +func (_c *MockGrpcClient_SetGetAddrFunc_Call[T]) RunAndReturn(run func(func() (string, error))) *MockGrpcClient_SetGetAddrFunc_Call[T] { + _c.Call.Return(run) + return _c +} + +// SetNewGrpcClientFunc provides a mock function with given fields: _a0 +func (_m *MockGrpcClient[T]) SetNewGrpcClientFunc(_a0 func(*grpc.ClientConn) T) { + _m.Called(_a0) +} + +// MockGrpcClient_SetNewGrpcClientFunc_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SetNewGrpcClientFunc' +type MockGrpcClient_SetNewGrpcClientFunc_Call[T grpcclient.GrpcComponent] struct { + *mock.Call +} + +// SetNewGrpcClientFunc is a helper method to define mock.On call +// - _a0 func(*grpc.ClientConn) T +func (_e *MockGrpcClient_Expecter[T]) SetNewGrpcClientFunc(_a0 interface{}) *MockGrpcClient_SetNewGrpcClientFunc_Call[T] { + return &MockGrpcClient_SetNewGrpcClientFunc_Call[T]{Call: _e.mock.On("SetNewGrpcClientFunc", _a0)} +} + +func (_c *MockGrpcClient_SetNewGrpcClientFunc_Call[T]) Run(run func(_a0 func(*grpc.ClientConn) T)) *MockGrpcClient_SetNewGrpcClientFunc_Call[T] { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(func(*grpc.ClientConn) T)) + }) + return _c +} + +func (_c *MockGrpcClient_SetNewGrpcClientFunc_Call[T]) Return() *MockGrpcClient_SetNewGrpcClientFunc_Call[T] { + _c.Call.Return() + return _c +} + +func (_c *MockGrpcClient_SetNewGrpcClientFunc_Call[T]) RunAndReturn(run func(func(*grpc.ClientConn) T)) *MockGrpcClient_SetNewGrpcClientFunc_Call[T] { + _c.Call.Return(run) + return _c +} + +// SetNodeID provides a mock function with given fields: _a0 +func (_m *MockGrpcClient[T]) SetNodeID(_a0 int64) { + _m.Called(_a0) +} + +// MockGrpcClient_SetNodeID_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SetNodeID' +type MockGrpcClient_SetNodeID_Call[T grpcclient.GrpcComponent] struct { + *mock.Call +} + +// SetNodeID is a helper method to define mock.On call +// - _a0 int64 +func (_e *MockGrpcClient_Expecter[T]) SetNodeID(_a0 interface{}) *MockGrpcClient_SetNodeID_Call[T] { + return &MockGrpcClient_SetNodeID_Call[T]{Call: _e.mock.On("SetNodeID", _a0)} +} + +func (_c *MockGrpcClient_SetNodeID_Call[T]) Run(run func(_a0 int64)) *MockGrpcClient_SetNodeID_Call[T] { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(int64)) + }) + return _c +} + +func (_c *MockGrpcClient_SetNodeID_Call[T]) Return() *MockGrpcClient_SetNodeID_Call[T] { + _c.Call.Return() + return _c +} + +func (_c *MockGrpcClient_SetNodeID_Call[T]) RunAndReturn(run func(int64)) *MockGrpcClient_SetNodeID_Call[T] { + _c.Call.Return(run) + return _c +} + +// SetRole provides a mock function with given fields: _a0 +func (_m *MockGrpcClient[T]) SetRole(_a0 string) { + _m.Called(_a0) +} + +// MockGrpcClient_SetRole_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SetRole' +type MockGrpcClient_SetRole_Call[T grpcclient.GrpcComponent] struct { + *mock.Call +} + +// SetRole is a helper method to define mock.On call +// - _a0 string +func (_e *MockGrpcClient_Expecter[T]) SetRole(_a0 interface{}) *MockGrpcClient_SetRole_Call[T] { + return &MockGrpcClient_SetRole_Call[T]{Call: _e.mock.On("SetRole", _a0)} +} + +func (_c *MockGrpcClient_SetRole_Call[T]) Run(run func(_a0 string)) *MockGrpcClient_SetRole_Call[T] { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(string)) + }) + return _c +} + +func (_c *MockGrpcClient_SetRole_Call[T]) Return() *MockGrpcClient_SetRole_Call[T] { + _c.Call.Return() + return _c +} + +func (_c *MockGrpcClient_SetRole_Call[T]) RunAndReturn(run func(string)) *MockGrpcClient_SetRole_Call[T] { + _c.Call.Return(run) + return _c +} + +// SetSession provides a mock function with given fields: sess +func (_m *MockGrpcClient[T]) SetSession(sess *sessionutil.Session) { + _m.Called(sess) +} + +// MockGrpcClient_SetSession_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SetSession' +type MockGrpcClient_SetSession_Call[T grpcclient.GrpcComponent] struct { + *mock.Call +} + +// SetSession is a helper method to define mock.On call +// - sess *sessionutil.Session +func (_e *MockGrpcClient_Expecter[T]) SetSession(sess interface{}) *MockGrpcClient_SetSession_Call[T] { + return &MockGrpcClient_SetSession_Call[T]{Call: _e.mock.On("SetSession", sess)} +} + +func (_c *MockGrpcClient_SetSession_Call[T]) Run(run func(sess *sessionutil.Session)) *MockGrpcClient_SetSession_Call[T] { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(*sessionutil.Session)) + }) + return _c +} + +func (_c *MockGrpcClient_SetSession_Call[T]) Return() *MockGrpcClient_SetSession_Call[T] { + _c.Call.Return() + return _c +} + +func (_c *MockGrpcClient_SetSession_Call[T]) RunAndReturn(run func(*sessionutil.Session)) *MockGrpcClient_SetSession_Call[T] { + _c.Call.Return(run) + return _c +} + +// NewMockGrpcClient creates a new instance of MockGrpcClient. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewMockGrpcClient[T grpcclient.GrpcComponent](t interface { + mock.TestingT + Cleanup(func()) +}) *MockGrpcClient[T] { + mock := &MockGrpcClient[T]{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/internal/rootcoord/broker_test.go b/internal/rootcoord/broker_test.go index 7d2788b3f6..a5a5ea00da 100644 --- a/internal/rootcoord/broker_test.go +++ b/internal/rootcoord/broker_test.go @@ -227,18 +227,19 @@ func TestServerBroker_GetSegmentIndexState(t *testing.T) { t.Run("success", func(t *testing.T) { c := newTestCore(withValidDataCoord()) - c.dataCoord.(*mockDataCoord).GetSegmentIndexStateFunc = func(ctx context.Context, req *indexpb.GetSegmentIndexStateRequest) (*indexpb.GetSegmentIndexStateResponse, error) { - return &indexpb.GetSegmentIndexStateResponse{ - Status: merr.Success(), - States: []*indexpb.SegmentIndexState{ - { - SegmentID: 1, - State: commonpb.IndexState_Finished, - FailReason: "", - }, + mockDataCoord := mocks.NewMockDataCoordClient(t) + mockDataCoord.EXPECT().GetSegmentIndexState(mock.Anything, mock.Anything).Return(&indexpb.GetSegmentIndexStateResponse{ + Status: merr.Success(), + States: []*indexpb.SegmentIndexState{ + { + SegmentID: 1, + State: commonpb.IndexState_Finished, + FailReason: "", }, - }, nil - } + }, + }, nil) + c.dataCoord = mockDataCoord + b := newServerBroker(c) ctx := context.Background() states, err := b.GetSegmentIndexState(ctx, 1, "index_name", []UniqueID{1}) diff --git a/internal/rootcoord/mock_test.go b/internal/rootcoord/mock_test.go index 0922debfe4..679189ff1d 100644 --- a/internal/rootcoord/mock_test.go +++ b/internal/rootcoord/mock_test.go @@ -238,29 +238,6 @@ func newMockMetaTable() *mockMetaTable { return &mockMetaTable{} } -//type mockIndexCoord struct { -// types.IndexCoord -// GetComponentStatesFunc func(ctx context.Context) (*milvuspb.ComponentStates, error) -// GetSegmentIndexStateFunc func(ctx context.Context, req *indexpb.GetSegmentIndexStateRequest) (*indexpb.GetSegmentIndexStateResponse, error) -// DropIndexFunc func(ctx context.Context, req *indexpb.DropIndexRequest) (*commonpb.Status, error) -//} -// -//func newMockIndexCoord() *mockIndexCoord { -// return &mockIndexCoord{} -//} -// -//func (m mockIndexCoord) GetComponentStates(ctx context.Context) (*milvuspb.ComponentStates, error) { -// return m.GetComponentStatesFunc(ctx) -//} -// -//func (m mockIndexCoord) GetSegmentIndexState(ctx context.Context, req *indexpb.GetSegmentIndexStateRequest) (*indexpb.GetSegmentIndexStateResponse, error) { -// return m.GetSegmentIndexStateFunc(ctx, req) -//} -// -//func (m mockIndexCoord) DropIndex(ctx context.Context, req *indexpb.DropIndexRequest) (*commonpb.Status, error) { -// return m.DropIndexFunc(ctx, req) -//} - type mockDataCoord struct { types.DataCoordClient GetComponentStatesFunc func(ctx context.Context) (*milvuspb.ComponentStates, error) diff --git a/internal/rootcoord/quota_center_test.go b/internal/rootcoord/quota_center_test.go index 01fb5d36d4..e14b6fb3e6 100644 --- a/internal/rootcoord/quota_center_test.go +++ b/internal/rootcoord/quota_center_test.go @@ -26,7 +26,6 @@ import ( "github.com/cockroachdb/errors" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/mock" - "google.golang.org/grpc" "github.com/milvus-io/milvus-proto/go-api/v2/commonpb" "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb" @@ -44,27 +43,6 @@ import ( "github.com/milvus-io/milvus/pkg/util/typeutil" ) -type dataCoordMockForQuota struct { - mockDataCoord - retErr bool - retFailStatus bool -} - -func (d *dataCoordMockForQuota) GetMetrics(ctx context.Context, request *milvuspb.GetMetricsRequest, opts ...grpc.CallOption) (*milvuspb.GetMetricsResponse, error) { - mockErr := errors.New("mock error") - if d.retErr { - return nil, mockErr - } - if d.retFailStatus { - return &milvuspb.GetMetricsResponse{ - Status: merr.Status(mockErr), - }, nil - } - return &milvuspb.GetMetricsResponse{ - Status: merr.Success(), - }, nil -} - func TestQuotaCenter(t *testing.T) { paramtable.Init() ctx, cancel := context.WithCancel(context.Background()) @@ -75,11 +53,15 @@ func TestQuotaCenter(t *testing.T) { pcm := newProxyClientManager(core.proxyCreator) + dc := mocks.NewMockDataCoordClient(t) + dc.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(nil, nil).Maybe() + t.Run("test QuotaCenter", func(t *testing.T) { qc := mocks.NewMockQueryCoordClient(t) meta := mockrootcoord.NewIMetaTable(t) + meta.EXPECT().GetCollectionByID(mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil, merr.ErrCollectionNotFound).Maybe() - quotaCenter := NewQuotaCenter(pcm, qc, &dataCoordMockForQuota{}, core.tsoAllocator, meta) + quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) go quotaCenter.run() time.Sleep(10 * time.Millisecond) quotaCenter.stop() @@ -90,27 +72,34 @@ func TestQuotaCenter(t *testing.T) { meta := mockrootcoord.NewIMetaTable(t) meta.EXPECT().GetCollectionByID(mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil, merr.ErrCollectionNotFound).Maybe() qc.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(&milvuspb.GetMetricsResponse{Status: merr.Success()}, nil) - quotaCenter := NewQuotaCenter(pcm, qc, &dataCoordMockForQuota{}, core.tsoAllocator, meta) + quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) err = quotaCenter.syncMetrics() assert.Error(t, err) // for empty response - quotaCenter = NewQuotaCenter(pcm, qc, &dataCoordMockForQuota{}, core.tsoAllocator, meta) + quotaCenter = NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) err = quotaCenter.syncMetrics() assert.Error(t, err) - quotaCenter = NewQuotaCenter(pcm, qc, &dataCoordMockForQuota{retFailStatus: true}, core.tsoAllocator, meta) + dc.ExpectedCalls = nil + dc.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(&milvuspb.GetMetricsResponse{ + Status: merr.Status(errors.New("mock error")), + }, nil) + + quotaCenter = NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) err = quotaCenter.syncMetrics() assert.Error(t, err) + dc.ExpectedCalls = nil + dc.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(nil, errors.New("mock error")) qc.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(nil, fmt.Errorf("mock err")) - quotaCenter = NewQuotaCenter(pcm, qc, &dataCoordMockForQuota{retErr: true}, core.tsoAllocator, meta) + quotaCenter = NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) err = quotaCenter.syncMetrics() assert.Error(t, err) qc.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(&milvuspb.GetMetricsResponse{ Status: merr.Status(err), }, nil) - quotaCenter = NewQuotaCenter(pcm, qc, &dataCoordMockForQuota{}, core.tsoAllocator, meta) + quotaCenter = NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) err = quotaCenter.syncMetrics() assert.Error(t, err) }) @@ -119,7 +108,7 @@ func TestQuotaCenter(t *testing.T) { qc := mocks.NewMockQueryCoordClient(t) meta := mockrootcoord.NewIMetaTable(t) meta.EXPECT().GetCollectionByID(mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil, merr.ErrCollectionNotFound).Maybe() - quotaCenter := NewQuotaCenter(pcm, qc, &dataCoordMockForQuota{}, core.tsoAllocator, meta) + quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) quotaCenter.readableCollections = []int64{1, 2, 3} quotaCenter.resetAllCurrentRates() quotaCenter.forceDenyReading(commonpb.ErrorCode_ForceDeny, 1, 2, 3, 4) @@ -149,7 +138,7 @@ func TestQuotaCenter(t *testing.T) { qc := mocks.NewMockQueryCoordClient(t) meta := mockrootcoord.NewIMetaTable(t) meta.EXPECT().GetCollectionByID(mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil, merr.ErrCollectionNotFound).Maybe() - quotaCenter := NewQuotaCenter(pcm, qc, &dataCoordMockForQuota{}, core.tsoAllocator, meta) + quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) err = quotaCenter.calculateRates() assert.NoError(t, err) alloc := newMockTsoAllocator() @@ -165,7 +154,7 @@ func TestQuotaCenter(t *testing.T) { qc := mocks.NewMockQueryCoordClient(t) meta := mockrootcoord.NewIMetaTable(t) meta.EXPECT().GetCollectionByID(mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil, merr.ErrCollectionNotFound).Maybe() - quotaCenter := NewQuotaCenter(pcm, qc, &dataCoordMockForQuota{}, core.tsoAllocator, meta) + quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) type ttCase struct { maxTtDelay time.Duration curTt time.Time @@ -213,7 +202,7 @@ func TestQuotaCenter(t *testing.T) { qc := mocks.NewMockQueryCoordClient(t) meta := mockrootcoord.NewIMetaTable(t) meta.EXPECT().GetCollectionByID(mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil, merr.ErrCollectionNotFound).Maybe() - quotaCenter := NewQuotaCenter(pcm, qc, &dataCoordMockForQuota{}, core.tsoAllocator, meta) + quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) type ttCase struct { delay time.Duration expectedFactor float64 @@ -286,7 +275,7 @@ func TestQuotaCenter(t *testing.T) { qc := mocks.NewMockQueryCoordClient(t) meta := mockrootcoord.NewIMetaTable(t) meta.EXPECT().GetCollectionByID(mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil, merr.ErrCollectionNotFound).Maybe() - quotaCenter := NewQuotaCenter(pcm, qc, &dataCoordMockForQuota{}, core.tsoAllocator, meta) + quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) quotaCenter.readableCollections = []int64{1, 2, 3} quotaCenter.proxyMetrics = map[UniqueID]*metricsinfo.ProxyQuotaMetrics{ 1: {Rms: []metricsinfo.RateMetric{ @@ -349,7 +338,7 @@ func TestQuotaCenter(t *testing.T) { qc := mocks.NewMockQueryCoordClient(t) meta := mockrootcoord.NewIMetaTable(t) meta.EXPECT().GetCollectionByID(mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil, merr.ErrCollectionNotFound).Maybe() - quotaCenter := NewQuotaCenter(pcm, qc, &dataCoordMockForQuota{}, core.tsoAllocator, meta) + quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) err = quotaCenter.calculateWriteRates() assert.NoError(t, err) @@ -389,7 +378,7 @@ func TestQuotaCenter(t *testing.T) { qc := mocks.NewMockQueryCoordClient(t) meta := mockrootcoord.NewIMetaTable(t) meta.EXPECT().GetCollectionByID(mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil, merr.ErrCollectionNotFound).Maybe() - quotaCenter := NewQuotaCenter(pcm, qc, &dataCoordMockForQuota{}, core.tsoAllocator, meta) + quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) type memCase struct { lowWater float64 highWater float64 @@ -444,7 +433,7 @@ func TestQuotaCenter(t *testing.T) { qc := mocks.NewMockQueryCoordClient(t) meta := mockrootcoord.NewIMetaTable(t) meta.EXPECT().GetCollectionByID(mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil, merr.ErrCollectionNotFound).Maybe() - quotaCenter := NewQuotaCenter(pcm, qc, &dataCoordMockForQuota{}, core.tsoAllocator, meta) + quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) defaultRatio := Params.QuotaConfig.GrowingSegmentsSizeMinRateRatio.GetAsFloat() tests := []struct { low float64 @@ -499,7 +488,7 @@ func TestQuotaCenter(t *testing.T) { qc := mocks.NewMockQueryCoordClient(t) meta := mockrootcoord.NewIMetaTable(t) meta.EXPECT().GetCollectionByID(mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil, merr.ErrCollectionNotFound).Maybe() - quotaCenter := NewQuotaCenter(pcm, qc, &dataCoordMockForQuota{}, core.tsoAllocator, meta) + quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) quotaCenter.checkDiskQuota() // total DiskQuota exceeded @@ -549,7 +538,7 @@ func TestQuotaCenter(t *testing.T) { }} meta := mockrootcoord.NewIMetaTable(t) meta.EXPECT().GetCollectionByID(mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil, merr.ErrCollectionNotFound).Maybe() - quotaCenter := NewQuotaCenter(pcm, qc, &dataCoordMockForQuota{}, core.tsoAllocator, meta) + quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) quotaCenter.resetAllCurrentRates() collectionID := int64(1) quotaCenter.currentRates[collectionID] = make(map[internalpb.RateType]ratelimitutil.Limit) @@ -565,7 +554,7 @@ func TestQuotaCenter(t *testing.T) { qc := mocks.NewMockQueryCoordClient(t) meta := mockrootcoord.NewIMetaTable(t) meta.EXPECT().GetCollectionByID(mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil, merr.ErrCollectionNotFound).Maybe() - quotaCenter := NewQuotaCenter(pcm, qc, &dataCoordMockForQuota{}, core.tsoAllocator, meta) + quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) collectionID := int64(1) quotaCenter.quotaStates[collectionID] = make(map[milvuspb.QuotaState]commonpb.ErrorCode) quotaCenter.quotaStates[collectionID][milvuspb.QuotaState_DenyToWrite] = commonpb.ErrorCode_MemoryQuotaExhausted @@ -577,7 +566,7 @@ func TestQuotaCenter(t *testing.T) { qc := mocks.NewMockQueryCoordClient(t) meta := mockrootcoord.NewIMetaTable(t) meta.EXPECT().GetCollectionByID(mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil, merr.ErrCollectionNotFound).Maybe() - quotaCenter := NewQuotaCenter(pcm, qc, &dataCoordMockForQuota{}, core.tsoAllocator, meta) + quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) quotaCenter.resetAllCurrentRates() minRate := Limit(100) collectionID := int64(1) @@ -606,7 +595,7 @@ func TestQuotaCenter(t *testing.T) { collection := UniqueID(0) meta := mockrootcoord.NewIMetaTable(t) meta.EXPECT().GetCollectionByID(mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil, merr.ErrCollectionNotFound).Maybe() - quotaCenter := NewQuotaCenter(pcm, nil, &dataCoordMockForQuota{}, core.tsoAllocator, meta) + quotaCenter := NewQuotaCenter(pcm, nil, dc, core.tsoAllocator, meta) quotaCenter.resetAllCurrentRates() quotaBackup := Params.QuotaConfig.DiskQuota.GetValue() colQuotaBackup := Params.QuotaConfig.DiskQuotaPerCollection.GetValue() @@ -628,7 +617,7 @@ func TestQuotaCenter(t *testing.T) { t.Run("test reset current rates", func(t *testing.T) { meta := mockrootcoord.NewIMetaTable(t) meta.EXPECT().GetCollectionByID(mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil, merr.ErrCollectionNotFound).Maybe() - quotaCenter := NewQuotaCenter(pcm, nil, &dataCoordMockForQuota{}, core.tsoAllocator, meta) + quotaCenter := NewQuotaCenter(pcm, nil, dc, core.tsoAllocator, meta) quotaCenter.readableCollections = []int64{1} quotaCenter.writableCollections = []int64{1} quotaCenter.resetAllCurrentRates() diff --git a/internal/rootcoord/root_coord_test.go b/internal/rootcoord/root_coord_test.go index ef5ec8497c..1086879cf5 100644 --- a/internal/rootcoord/root_coord_test.go +++ b/internal/rootcoord/root_coord_test.go @@ -37,6 +37,7 @@ import ( memkv "github.com/milvus-io/milvus/internal/kv/mem" "github.com/milvus-io/milvus/internal/kv/mocks" "github.com/milvus-io/milvus/internal/metastore/model" + mocksutil "github.com/milvus-io/milvus/internal/mocks" "github.com/milvus-io/milvus/internal/proto/datapb" "github.com/milvus-io/milvus/internal/proto/etcdpb" "github.com/milvus-io/milvus/internal/proto/internalpb" @@ -1444,23 +1445,8 @@ func TestCore_ReportImport(t *testing.T) { return nil } - dc := newMockDataCoord() - dc.GetComponentStatesFunc = func(ctx context.Context) (*milvuspb.ComponentStates, error) { - return &milvuspb.ComponentStates{ - State: &milvuspb.ComponentInfo{ - NodeID: TestRootCoordID, - StateCode: commonpb.StateCode_Healthy, - }, - SubcomponentStates: nil, - Status: merr.Success(), - }, nil - } - dc.WatchChannelsFunc = func(ctx context.Context, req *datapb.WatchChannelsRequest) (*datapb.WatchChannelsResponse, error) { - return &datapb.WatchChannelsResponse{Status: merr.Success()}, nil - } - dc.FlushFunc = func(ctx context.Context, req *datapb.FlushRequest) (*datapb.FlushResponse, error) { - return &datapb.FlushResponse{Status: merr.Success()}, nil - } + dc := mocksutil.NewMockDataCoordClient(t) + dc.EXPECT().Flush(mock.Anything, mock.Anything).Return(&datapb.FlushResponse{Status: merr.Success()}, nil) mockCallImportServiceErr := false callImportServiceFn := func(ctx context.Context, req *datapb.ImportTaskRequest) (*datapb.ImportTaskResponse, error) { diff --git a/internal/util/mock/grpc_datacoord_client.go b/internal/util/mock/grpc_datacoord_client.go deleted file mode 100644 index 7afd848a39..0000000000 --- a/internal/util/mock/grpc_datacoord_client.go +++ /dev/null @@ -1,236 +0,0 @@ -// Licensed to the LF AI & Data foundation under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package mock - -import ( - "context" - - "google.golang.org/grpc" - - "github.com/milvus-io/milvus-proto/go-api/v2/commonpb" - "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb" - "github.com/milvus-io/milvus/internal/proto/datapb" - "github.com/milvus-io/milvus/internal/proto/indexpb" - "github.com/milvus-io/milvus/internal/proto/internalpb" - "github.com/milvus-io/milvus/pkg/util/merr" - "github.com/milvus-io/milvus/pkg/util/uniquegenerator" -) - -var _ datapb.DataCoordClient = &GrpcDataCoordClient{} - -// GrpcDataCoordClient mocks of GrpcDataCoordClient -type GrpcDataCoordClient struct { - Err error -} - -func (m *GrpcDataCoordClient) GcConfirm(ctx context.Context, in *datapb.GcConfirmRequest, opts ...grpc.CallOption) (*datapb.GcConfirmResponse, error) { - return &datapb.GcConfirmResponse{}, m.Err -} - -func (m *GrpcDataCoordClient) CheckHealth(ctx context.Context, in *milvuspb.CheckHealthRequest, opts ...grpc.CallOption) (*milvuspb.CheckHealthResponse, error) { - return &milvuspb.CheckHealthResponse{}, m.Err -} - -func (m *GrpcDataCoordClient) GetComponentStates(ctx context.Context, in *milvuspb.GetComponentStatesRequest, opts ...grpc.CallOption) (*milvuspb.ComponentStates, error) { - return &milvuspb.ComponentStates{ - State: &milvuspb.ComponentInfo{ - NodeID: int64(uniquegenerator.GetUniqueIntGeneratorIns().GetInt()), - Role: "MockDataCoord", - StateCode: commonpb.StateCode_Healthy, - ExtraInfo: nil, - }, - SubcomponentStates: nil, - Status: merr.Success(), - }, m.Err -} - -func (m *GrpcDataCoordClient) GetTimeTickChannel(ctx context.Context, in *internalpb.GetTimeTickChannelRequest, opts ...grpc.CallOption) (*milvuspb.StringResponse, error) { - return &milvuspb.StringResponse{}, m.Err -} - -func (m *GrpcDataCoordClient) GetStatisticsChannel(ctx context.Context, in *internalpb.GetStatisticsChannelRequest, opts ...grpc.CallOption) (*milvuspb.StringResponse, error) { - return &milvuspb.StringResponse{}, m.Err -} - -func (m *GrpcDataCoordClient) Flush(ctx context.Context, in *datapb.FlushRequest, opts ...grpc.CallOption) (*datapb.FlushResponse, error) { - return &datapb.FlushResponse{}, m.Err -} - -func (m *GrpcDataCoordClient) AssignSegmentID(ctx context.Context, in *datapb.AssignSegmentIDRequest, opts ...grpc.CallOption) (*datapb.AssignSegmentIDResponse, error) { - return &datapb.AssignSegmentIDResponse{}, m.Err -} - -func (m *GrpcDataCoordClient) GetSegmentInfo(ctx context.Context, in *datapb.GetSegmentInfoRequest, opts ...grpc.CallOption) (*datapb.GetSegmentInfoResponse, error) { - return &datapb.GetSegmentInfoResponse{}, m.Err -} - -func (m *GrpcDataCoordClient) GetSegmentStates(ctx context.Context, in *datapb.GetSegmentStatesRequest, opts ...grpc.CallOption) (*datapb.GetSegmentStatesResponse, error) { - return &datapb.GetSegmentStatesResponse{}, m.Err -} - -func (m *GrpcDataCoordClient) GetInsertBinlogPaths(ctx context.Context, in *datapb.GetInsertBinlogPathsRequest, opts ...grpc.CallOption) (*datapb.GetInsertBinlogPathsResponse, error) { - return &datapb.GetInsertBinlogPathsResponse{}, m.Err -} - -func (m *GrpcDataCoordClient) GetCollectionStatistics(ctx context.Context, in *datapb.GetCollectionStatisticsRequest, opts ...grpc.CallOption) (*datapb.GetCollectionStatisticsResponse, error) { - return &datapb.GetCollectionStatisticsResponse{}, m.Err -} - -func (m *GrpcDataCoordClient) GetPartitionStatistics(ctx context.Context, in *datapb.GetPartitionStatisticsRequest, opts ...grpc.CallOption) (*datapb.GetPartitionStatisticsResponse, error) { - return &datapb.GetPartitionStatisticsResponse{}, m.Err -} - -func (m *GrpcDataCoordClient) GetSegmentInfoChannel(ctx context.Context, in *datapb.GetSegmentInfoChannelRequest, opts ...grpc.CallOption) (*milvuspb.StringResponse, error) { - return &milvuspb.StringResponse{}, m.Err -} - -func (m *GrpcDataCoordClient) SaveBinlogPaths(ctx context.Context, in *datapb.SaveBinlogPathsRequest, opts ...grpc.CallOption) (*commonpb.Status, error) { - return &commonpb.Status{}, m.Err -} - -func (m *GrpcDataCoordClient) GetRecoveryInfo(ctx context.Context, in *datapb.GetRecoveryInfoRequest, opts ...grpc.CallOption) (*datapb.GetRecoveryInfoResponse, error) { - return &datapb.GetRecoveryInfoResponse{}, m.Err -} - -func (m *GrpcDataCoordClient) GetRecoveryInfoV2(ctx context.Context, in *datapb.GetRecoveryInfoRequestV2, opts ...grpc.CallOption) (*datapb.GetRecoveryInfoResponseV2, error) { - return &datapb.GetRecoveryInfoResponseV2{}, m.Err -} - -func (m *GrpcDataCoordClient) GetFlushedSegments(ctx context.Context, in *datapb.GetFlushedSegmentsRequest, opts ...grpc.CallOption) (*datapb.GetFlushedSegmentsResponse, error) { - return &datapb.GetFlushedSegmentsResponse{}, m.Err -} - -func (m *GrpcDataCoordClient) GetSegmentsByStates(ctx context.Context, in *datapb.GetSegmentsByStatesRequest, opts ...grpc.CallOption) (*datapb.GetSegmentsByStatesResponse, error) { - return &datapb.GetSegmentsByStatesResponse{}, m.Err -} - -func (m *GrpcDataCoordClient) ShowConfigurations(ctx context.Context, in *internalpb.ShowConfigurationsRequest, opts ...grpc.CallOption) (*internalpb.ShowConfigurationsResponse, error) { - return &internalpb.ShowConfigurationsResponse{}, m.Err -} - -func (m *GrpcDataCoordClient) GetMetrics(ctx context.Context, in *milvuspb.GetMetricsRequest, opts ...grpc.CallOption) (*milvuspb.GetMetricsResponse, error) { - return &milvuspb.GetMetricsResponse{}, m.Err -} - -func (m *GrpcDataCoordClient) CompleteCompaction(ctx context.Context, req *datapb.CompactionPlanResult, opts ...grpc.CallOption) (*commonpb.Status, error) { - return &commonpb.Status{}, m.Err -} - -func (m *GrpcDataCoordClient) ManualCompaction(ctx context.Context, in *milvuspb.ManualCompactionRequest, opts ...grpc.CallOption) (*milvuspb.ManualCompactionResponse, error) { - return &milvuspb.ManualCompactionResponse{}, m.Err -} - -func (m *GrpcDataCoordClient) GetCompactionState(ctx context.Context, in *milvuspb.GetCompactionStateRequest, opts ...grpc.CallOption) (*milvuspb.GetCompactionStateResponse, error) { - return &milvuspb.GetCompactionStateResponse{}, m.Err -} - -func (m *GrpcDataCoordClient) GetCompactionStateWithPlans(ctx context.Context, req *milvuspb.GetCompactionPlansRequest, opts ...grpc.CallOption) (*milvuspb.GetCompactionPlansResponse, error) { - return &milvuspb.GetCompactionPlansResponse{}, m.Err -} - -func (m *GrpcDataCoordClient) WatchChannels(ctx context.Context, req *datapb.WatchChannelsRequest, opts ...grpc.CallOption) (*datapb.WatchChannelsResponse, error) { - return &datapb.WatchChannelsResponse{}, m.Err -} - -func (m *GrpcDataCoordClient) GetFlushState(ctx context.Context, req *datapb.GetFlushStateRequest, opts ...grpc.CallOption) (*milvuspb.GetFlushStateResponse, error) { - return &milvuspb.GetFlushStateResponse{}, m.Err -} - -func (m *GrpcDataCoordClient) GetFlushAllState(ctx context.Context, req *milvuspb.GetFlushAllStateRequest, opts ...grpc.CallOption) (*milvuspb.GetFlushAllStateResponse, error) { - return &milvuspb.GetFlushAllStateResponse{}, m.Err -} - -func (m *GrpcDataCoordClient) DropVirtualChannel(ctx context.Context, req *datapb.DropVirtualChannelRequest, opts ...grpc.CallOption) (*datapb.DropVirtualChannelResponse, error) { - return &datapb.DropVirtualChannelResponse{}, m.Err -} - -func (m *GrpcDataCoordClient) SetSegmentState(ctx context.Context, req *datapb.SetSegmentStateRequest, opts ...grpc.CallOption) (*datapb.SetSegmentStateResponse, error) { - return &datapb.SetSegmentStateResponse{}, m.Err -} - -func (m *GrpcDataCoordClient) Import(ctx context.Context, req *datapb.ImportTaskRequest, opts ...grpc.CallOption) (*datapb.ImportTaskResponse, error) { - return &datapb.ImportTaskResponse{}, m.Err -} - -func (m *GrpcDataCoordClient) UpdateSegmentStatistics(ctx context.Context, req *datapb.UpdateSegmentStatisticsRequest, opts ...grpc.CallOption) (*commonpb.Status, error) { - return &commonpb.Status{}, m.Err -} - -func (m *GrpcDataCoordClient) UpdateChannelCheckpoint(ctx context.Context, req *datapb.UpdateChannelCheckpointRequest, opts ...grpc.CallOption) (*commonpb.Status, error) { - return &commonpb.Status{}, m.Err -} - -func (m *GrpcDataCoordClient) SaveImportSegment(ctx context.Context, in *datapb.SaveImportSegmentRequest, opts ...grpc.CallOption) (*commonpb.Status, error) { - return &commonpb.Status{}, m.Err -} - -func (m *GrpcDataCoordClient) UnsetIsImportingState(context.Context, *datapb.UnsetIsImportingStateRequest, ...grpc.CallOption) (*commonpb.Status, error) { - return &commonpb.Status{}, m.Err -} - -func (m *GrpcDataCoordClient) MarkSegmentsDropped(context.Context, *datapb.MarkSegmentsDroppedRequest, ...grpc.CallOption) (*commonpb.Status, error) { - return &commonpb.Status{}, m.Err -} - -func (m *GrpcDataCoordClient) BroadcastAlteredCollection(ctx context.Context, in *datapb.AlterCollectionRequest, opts ...grpc.CallOption) (*commonpb.Status, error) { - return &commonpb.Status{}, m.Err -} - -func (m *GrpcDataCoordClient) CreateIndex(ctx context.Context, req *indexpb.CreateIndexRequest, opts ...grpc.CallOption) (*commonpb.Status, error) { - return &commonpb.Status{}, m.Err -} - -func (m *GrpcDataCoordClient) DropIndex(ctx context.Context, req *indexpb.DropIndexRequest, opts ...grpc.CallOption) (*commonpb.Status, error) { - return &commonpb.Status{}, m.Err -} - -func (m *GrpcDataCoordClient) GetIndexState(ctx context.Context, req *indexpb.GetIndexStateRequest, opts ...grpc.CallOption) (*indexpb.GetIndexStateResponse, error) { - return &indexpb.GetIndexStateResponse{}, m.Err -} - -// GetSegmentIndexState gets the index state of the segments in the request from RootCoord. -func (m *GrpcDataCoordClient) GetSegmentIndexState(ctx context.Context, req *indexpb.GetSegmentIndexStateRequest, opts ...grpc.CallOption) (*indexpb.GetSegmentIndexStateResponse, error) { - return &indexpb.GetSegmentIndexStateResponse{}, m.Err -} - -// GetIndexInfos gets the index files of the IndexBuildIDs in the request from RootCoordinator. -func (m *GrpcDataCoordClient) GetIndexInfos(ctx context.Context, req *indexpb.GetIndexInfoRequest, opts ...grpc.CallOption) (*indexpb.GetIndexInfoResponse, error) { - return &indexpb.GetIndexInfoResponse{}, m.Err -} - -// DescribeIndex describe the index info of the collection. -func (m *GrpcDataCoordClient) DescribeIndex(ctx context.Context, req *indexpb.DescribeIndexRequest, opts ...grpc.CallOption) (*indexpb.DescribeIndexResponse, error) { - return &indexpb.DescribeIndexResponse{}, m.Err -} - -// GetIndexStatistics get the information of index. -func (m *GrpcDataCoordClient) GetIndexStatistics(ctx context.Context, in *indexpb.GetIndexStatisticsRequest, opts ...grpc.CallOption) (*indexpb.GetIndexStatisticsResponse, error) { - return &indexpb.GetIndexStatisticsResponse{}, m.Err -} - -// GetIndexBuildProgress get the index building progress by num rows. -func (m *GrpcDataCoordClient) GetIndexBuildProgress(ctx context.Context, req *indexpb.GetIndexBuildProgressRequest, opts ...grpc.CallOption) (*indexpb.GetIndexBuildProgressResponse, error) { - return &indexpb.GetIndexBuildProgressResponse{}, m.Err -} - -func (m *GrpcDataCoordClient) ReportDataNodeTtMsgs(ctx context.Context, in *datapb.ReportDataNodeTtMsgsRequest, opts ...grpc.CallOption) (*commonpb.Status, error) { - return &commonpb.Status{}, m.Err -} - -func (m *GrpcDataCoordClient) Close() error { - return nil -} diff --git a/internal/util/mock/grpc_proxy_client.go b/internal/util/mock/grpc_proxy_client.go deleted file mode 100644 index 5d281d2efa..0000000000 --- a/internal/util/mock/grpc_proxy_client.go +++ /dev/null @@ -1,74 +0,0 @@ -// Licensed to the LF AI & Data foundation under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package mock - -import ( - "context" - - "google.golang.org/grpc" - - "github.com/milvus-io/milvus-proto/go-api/v2/commonpb" - "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb" - "github.com/milvus-io/milvus/internal/proto/internalpb" - "github.com/milvus-io/milvus/internal/proto/proxypb" -) - -var _ proxypb.ProxyClient = &GrpcProxyClient{} - -type GrpcProxyClient struct { - Err error -} - -func (m *GrpcProxyClient) RefreshPolicyInfoCache(ctx context.Context, in *proxypb.RefreshPolicyInfoCacheRequest, opts ...grpc.CallOption) (*commonpb.Status, error) { - return &commonpb.Status{}, m.Err -} - -func (m *GrpcProxyClient) GetComponentStates(ctx context.Context, in *milvuspb.GetComponentStatesRequest, opts ...grpc.CallOption) (*milvuspb.ComponentStates, error) { - return &milvuspb.ComponentStates{}, m.Err -} - -func (m *GrpcProxyClient) GetStatisticsChannel(ctx context.Context, in *internalpb.GetStatisticsChannelRequest, opts ...grpc.CallOption) (*milvuspb.StringResponse, error) { - return &milvuspb.StringResponse{}, m.Err -} - -func (m *GrpcProxyClient) InvalidateCollectionMetaCache(ctx context.Context, in *proxypb.InvalidateCollMetaCacheRequest, opts ...grpc.CallOption) (*commonpb.Status, error) { - return &commonpb.Status{}, m.Err -} - -func (m *GrpcProxyClient) GetDdChannel(ctx context.Context, in *internalpb.GetDdChannelRequest, opts ...grpc.CallOption) (*milvuspb.StringResponse, error) { - return &milvuspb.StringResponse{}, m.Err -} - -func (m *GrpcProxyClient) InvalidateCredentialCache(ctx context.Context, in *proxypb.InvalidateCredCacheRequest, opts ...grpc.CallOption) (*commonpb.Status, error) { - return &commonpb.Status{}, m.Err -} - -func (m *GrpcProxyClient) UpdateCredentialCache(ctx context.Context, in *proxypb.UpdateCredCacheRequest, opts ...grpc.CallOption) (*commonpb.Status, error) { - return &commonpb.Status{}, m.Err -} - -func (m *GrpcProxyClient) GetProxyMetrics(ctx context.Context, in *milvuspb.GetMetricsRequest, opts ...grpc.CallOption) (*milvuspb.GetMetricsResponse, error) { - return &milvuspb.GetMetricsResponse{}, m.Err -} - -func (m *GrpcProxyClient) SetRates(ctx context.Context, in *proxypb.SetRatesRequest, opts ...grpc.CallOption) (*commonpb.Status, error) { - return &commonpb.Status{}, m.Err -} - -func (m *GrpcProxyClient) ListClientInfos(ctx context.Context, in *proxypb.ListClientInfosRequest, opts ...grpc.CallOption) (*proxypb.ListClientInfosResponse, error) { - return &proxypb.ListClientInfosResponse{}, m.Err -}