enhance: Use mockery to replace manual mock code (#29074)

issue: #29043
This PR remove mannul mock code for proxy and data coord

---------

Signed-off-by: Wei Liu <wei.liu@zilliz.com>

Makefile

@@ -455,6 +455,7 @@ 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

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)
 

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)
+}

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{
+		mdc := mocks.NewMockDataCoordClient(t)
-			stateErr: commonpb.ErrorCode_Success,
+		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{
+		mrc := mocks.NewMockRootCoordClient(t)
-			stateErr: commonpb.ErrorCode_Success,
+		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()

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() })
 

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
+}

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) {
+		mockDataCoord := mocks.NewMockDataCoordClient(t)
-			return &indexpb.GetSegmentIndexStateResponse{
+		mockDataCoord.EXPECT().GetSegmentIndexState(mock.Anything, mock.Anything).Return(&indexpb.GetSegmentIndexStateResponse{
-				Status: merr.Success(),
+			Status: merr.Success(),
-				States: []*indexpb.SegmentIndexState{
+			States: []*indexpb.SegmentIndexState{
-					{
+				{
-						SegmentID:  1,
+					SegmentID:  1,
-						State:      commonpb.IndexState_Finished,
+					State:      commonpb.IndexState_Finished,
-						FailReason: "",
+					FailReason: "",
-					},
 				},
-			}, nil
+			},
-		}
+		}, nil)
+		c.dataCoord = mockDataCoord
+
 		b := newServerBroker(c)
 		ctx := context.Background()
 		states, err := b.GetSegmentIndexState(ctx, 1, "index_name", []UniqueID{1})

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)

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()

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 := mocksutil.NewMockDataCoordClient(t)
-	dc.GetComponentStatesFunc = func(ctx context.Context) (*milvuspb.ComponentStates, error) {
+	dc.EXPECT().Flush(mock.Anything, mock.Anything).Return(&datapb.FlushResponse{Status: merr.Success()}, nil)
-		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
-	}
 
 	mockCallImportServiceErr := false
 	callImportServiceFn := func(ctx context.Context, req *datapb.ImportTaskRequest) (*datapb.ImportTaskResponse, error) {

internal/util/mock/grpc_datacoord_client.go

@@ -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
-}

internal/util/mock/grpc_proxy_client.go

@@ -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
-}