test/milvus
1
0
Fork 0
mirror of https://gitee.com/milvus-io/milvus.git synced 2024-11-29 18:38:44 +08:00
Code Issues Actions 11 Packages Projects Releases Wiki Activity

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

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

---------

Signed-off-by: Wei Liu <wei.liu@zilliz.com>
This commit is contained in:
wei liu 2023-12-13 10:46:44 +08:00 committed by GitHub
parent 990d723496
commit fe1eeae2aa
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
15 changed files with 3396 additions and 1995 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
Makefile
View File

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

1945
internal/distributed/datacoord/client/client_test.go
View File

File diff suppressed because it is too large Load Diff

1046
internal/distributed/datacoord/service_test.go
View File

File diff suppressed because it is too large Load Diff

118
internal/distributed/datanode/service_test.go
View File

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

488
internal/distributed/proxy/client/client_test.go
View File

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

750
internal/distributed/proxy/service_test.go
View File

File diff suppressed because it is too large Load Diff

54
internal/distributed/querycoord/service_test.go
View File

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

66
internal/distributed/rootcoord/service_test.go
View File

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

472
internal/mocks/mock_grpc_client.go Normal file
View File

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

23
internal/rootcoord/broker_test.go
View File

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

23
internal/rootcoord/mock_test.go
View File

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

73
internal/rootcoord/quota_center_test.go
View File

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

20
internal/rootcoord/root_coord_test.go
View File

@ -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) {

236
internal/util/mock/grpc_datacoord_client.go
View File

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

74
internal/util/mock/grpc_proxy_client.go
View File

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