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

1
Makefile
View File

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

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

File diff suppressed because it is too large Load Diff

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

@ -18,673 +18,319 @@ package grpcdatacoord
import (
"context"
"fmt"
"testing"
"github.com/cockroachdb/errors"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/tikv/client-go/v2/txnkv"
clientv3 "go.etcd.io/etcd/client/v3"
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
"github.com/milvus-io/milvus/internal/mocks"
"github.com/milvus-io/milvus/internal/proto/datapb"
"github.com/milvus-io/milvus/internal/proto/indexpb"
"github.com/milvus-io/milvus/internal/proto/internalpb"
"github.com/milvus-io/milvus/internal/types"
"github.com/milvus-io/milvus/pkg/util/merr"
"github.com/milvus-io/milvus/pkg/util/paramtable"
"github.com/milvus-io/milvus/pkg/util/tikv"
)
type MockDataCoord struct {
types.DataCoord
states *milvuspb.ComponentStates
status *commonpb.Status
err error
initErr error
startErr error
stopErr error
regErr error
strResp *milvuspb.StringResponse
infoResp *datapb.GetSegmentInfoResponse
flushResp *datapb.FlushResponse
assignResp *datapb.AssignSegmentIDResponse
segStateResp *datapb.GetSegmentStatesResponse
binResp *datapb.GetInsertBinlogPathsResponse
colStatResp *datapb.GetCollectionStatisticsResponse
partStatResp *datapb.GetPartitionStatisticsResponse
recoverResp *datapb.GetRecoveryInfoResponse
flushSegResp *datapb.GetFlushedSegmentsResponse
SegByStatesResp *datapb.GetSegmentsByStatesResponse
configResp *internalpb.ShowConfigurationsResponse
metricResp *milvuspb.GetMetricsResponse
compactionStateResp *milvuspb.GetCompactionStateResponse
manualCompactionResp *milvuspb.ManualCompactionResponse
compactionPlansResp *milvuspb.GetCompactionPlansResponse
watchChannelsResp *datapb.WatchChannelsResponse
getFlushStateResp *milvuspb.GetFlushStateResponse
getFlushAllStateResp *milvuspb.GetFlushAllStateResponse
dropVChanResp *datapb.DropVirtualChannelResponse
setSegmentStateResp *datapb.SetSegmentStateResponse
importResp *datapb.ImportTaskResponse
updateSegStatResp *commonpb.Status
updateChanPos *commonpb.Status
addSegmentResp *commonpb.Status
unsetIsImportingStateResp *commonpb.Status
markSegmentsDroppedResp *commonpb.Status
broadCastResp *commonpb.Status
createIndexResp *commonpb.Status
describeIndexResp *indexpb.DescribeIndexResponse
getIndexStatisticsResp *indexpb.GetIndexStatisticsResponse
dropIndexResp *commonpb.Status
getIndexStateResp *indexpb.GetIndexStateResponse
getIndexBuildProgressResp *indexpb.GetIndexBuildProgressResponse
getSegmentIndexStateResp *indexpb.GetSegmentIndexStateResponse
getIndexInfosResp *indexpb.GetIndexInfoResponse
}
func (m *MockDataCoord) Init() error {
return m.initErr
}
func (m *MockDataCoord) Start() error {
return m.startErr
}
func (m *MockDataCoord) Stop() error {
return m.stopErr
}
func (m *MockDataCoord) Register() error {
return m.regErr
}
func (*MockDataCoord) SetAddress(address string) {
}
func (m *MockDataCoord) SetEtcdClient(etcdClient *clientv3.Client) {
}
func (m *MockDataCoord) SetTiKVClient(client *txnkv.Client) {
}
func (m *MockDataCoord) SetRootCoordClient(rootCoord types.RootCoordClient) {
}
func (m *MockDataCoord) SetDataNodeCreator(func(context.Context, string, int64) (types.DataNodeClient, error)) {
}
func (m *MockDataCoord) SetIndexNodeCreator(func(context.Context, string, int64) (types.IndexNodeClient, error)) {
}
func (m *MockDataCoord) GetComponentStates(ctx context.Context, req *milvuspb.GetComponentStatesRequest) (*milvuspb.ComponentStates, error) {
return m.states, m.err
}
func (m *MockDataCoord) GetTimeTickChannel(ctx context.Context, req *internalpb.GetTimeTickChannelRequest) (*milvuspb.StringResponse, error) {
return m.strResp, m.err
}
func (m *MockDataCoord) GetStatisticsChannel(ctx context.Context, req *internalpb.GetStatisticsChannelRequest) (*milvuspb.StringResponse, error) {
return m.strResp, m.err
}
func (m *MockDataCoord) GetSegmentInfo(ctx context.Context, req *datapb.GetSegmentInfoRequest) (*datapb.GetSegmentInfoResponse, error) {
return m.infoResp, m.err
}
func (m *MockDataCoord) Flush(ctx context.Context, req *datapb.FlushRequest) (*datapb.FlushResponse, error) {
return m.flushResp, m.err
}
func (m *MockDataCoord) AssignSegmentID(ctx context.Context, req *datapb.AssignSegmentIDRequest) (*datapb.AssignSegmentIDResponse, error) {
return m.assignResp, m.err
}
func (m *MockDataCoord) GetSegmentStates(ctx context.Context, req *datapb.GetSegmentStatesRequest) (*datapb.GetSegmentStatesResponse, error) {
return m.segStateResp, m.err
}
func (m *MockDataCoord) GetInsertBinlogPaths(ctx context.Context, req *datapb.GetInsertBinlogPathsRequest) (*datapb.GetInsertBinlogPathsResponse, error) {
return m.binResp, m.err
}
func (m *MockDataCoord) GetCollectionStatistics(ctx context.Context, req *datapb.GetCollectionStatisticsRequest) (*datapb.GetCollectionStatisticsResponse, error) {
return m.colStatResp, m.err
}
func (m *MockDataCoord) GetPartitionStatistics(ctx context.Context, req *datapb.GetPartitionStatisticsRequest) (*datapb.GetPartitionStatisticsResponse, error) {
return m.partStatResp, m.err
}
func (m *MockDataCoord) GetSegmentInfoChannel(ctx context.Context, req *datapb.GetSegmentInfoChannelRequest) (*milvuspb.StringResponse, error) {
return m.strResp, m.err
}
func (m *MockDataCoord) SaveBinlogPaths(ctx context.Context, req *datapb.SaveBinlogPathsRequest) (*commonpb.Status, error) {
return m.status, m.err
}
func (m *MockDataCoord) GetRecoveryInfo(ctx context.Context, req *datapb.GetRecoveryInfoRequest) (*datapb.GetRecoveryInfoResponse, error) {
return m.recoverResp, m.err
}
func (m *MockDataCoord) GetFlushedSegments(ctx context.Context, req *datapb.GetFlushedSegmentsRequest) (*datapb.GetFlushedSegmentsResponse, error) {
return m.flushSegResp, m.err
}
func (m *MockDataCoord) GetSegmentsByStates(ctx context.Context, req *datapb.GetSegmentsByStatesRequest) (*datapb.GetSegmentsByStatesResponse, error) {
return m.SegByStatesResp, m.err
}
func (m *MockDataCoord) ShowConfigurations(ctx context.Context, req *internalpb.ShowConfigurationsRequest) (*internalpb.ShowConfigurationsResponse, error) {
return m.configResp, m.err
}
func (m *MockDataCoord) GetMetrics(ctx context.Context, req *milvuspb.GetMetricsRequest) (*milvuspb.GetMetricsResponse, error) {
return m.metricResp, m.err
}
func (m *MockDataCoord) CompleteCompaction(ctx context.Context, req *datapb.CompactionPlanResult) (*commonpb.Status, error) {
return m.status, m.err
}
func (m *MockDataCoord) ManualCompaction(ctx context.Context, req *milvuspb.ManualCompactionRequest) (*milvuspb.ManualCompactionResponse, error) {
return m.manualCompactionResp, m.err
}
func (m *MockDataCoord) GetCompactionState(ctx context.Context, req *milvuspb.GetCompactionStateRequest) (*milvuspb.GetCompactionStateResponse, error) {
return m.compactionStateResp, m.err
}
func (m *MockDataCoord) GetCompactionStateWithPlans(ctx context.Context, req *milvuspb.GetCompactionPlansRequest) (*milvuspb.GetCompactionPlansResponse, error) {
return m.compactionPlansResp, m.err
}
func (m *MockDataCoord) WatchChannels(ctx context.Context, req *datapb.WatchChannelsRequest) (*datapb.WatchChannelsResponse, error) {
return m.watchChannelsResp, m.err
}
func (m *MockDataCoord) GetFlushState(ctx context.Context, req *datapb.GetFlushStateRequest) (*milvuspb.GetFlushStateResponse, error) {
return m.getFlushStateResp, m.err
}
func (m *MockDataCoord) GetFlushAllState(ctx context.Context, req *milvuspb.GetFlushAllStateRequest) (*milvuspb.GetFlushAllStateResponse, error) {
return m.getFlushAllStateResp, m.err
}
func (m *MockDataCoord) DropVirtualChannel(ctx context.Context, req *datapb.DropVirtualChannelRequest) (*datapb.DropVirtualChannelResponse, error) {
return m.dropVChanResp, m.err
}
func (m *MockDataCoord) SetSegmentState(ctx context.Context, req *datapb.SetSegmentStateRequest) (*datapb.SetSegmentStateResponse, error) {
return m.setSegmentStateResp, m.err
}
func (m *MockDataCoord) Import(ctx context.Context, req *datapb.ImportTaskRequest) (*datapb.ImportTaskResponse, error) {
return m.importResp, m.err
}
func (m *MockDataCoord) UpdateSegmentStatistics(ctx context.Context, req *datapb.UpdateSegmentStatisticsRequest) (*commonpb.Status, error) {
return m.updateSegStatResp, m.err
}
func (m *MockDataCoord) UpdateChannelCheckpoint(ctx context.Context, req *datapb.UpdateChannelCheckpointRequest) (*commonpb.Status, error) {
return m.updateChanPos, m.err
}
func (m *MockDataCoord) SaveImportSegment(ctx context.Context, req *datapb.SaveImportSegmentRequest) (*commonpb.Status, error) {
return m.addSegmentResp, m.err
}
func (m *MockDataCoord) UnsetIsImportingState(context.Context, *datapb.UnsetIsImportingStateRequest) (*commonpb.Status, error) {
return m.unsetIsImportingStateResp, m.err
}
func (m *MockDataCoord) MarkSegmentsDropped(ctx context.Context, req *datapb.MarkSegmentsDroppedRequest) (*commonpb.Status, error) {
return m.markSegmentsDroppedResp, m.err
}
func (m *MockDataCoord) BroadcastAlteredCollection(ctx context.Context, req *datapb.AlterCollectionRequest) (*commonpb.Status, error) {
return m.broadCastResp, m.err
}
func (m *MockDataCoord) CheckHealth(ctx context.Context, req *milvuspb.CheckHealthRequest) (*milvuspb.CheckHealthResponse, error) {
return &milvuspb.CheckHealthResponse{
IsHealthy: true,
}, nil
}
func (m *MockDataCoord) CreateIndex(ctx context.Context, req *indexpb.CreateIndexRequest) (*commonpb.Status, error) {
return m.createIndexResp, m.err
}
func (m *MockDataCoord) DescribeIndex(ctx context.Context, req *indexpb.DescribeIndexRequest) (*indexpb.DescribeIndexResponse, error) {
return m.describeIndexResp, m.err
}
func (m *MockDataCoord) GetIndexStatistics(ctx context.Context, req *indexpb.GetIndexStatisticsRequest) (*indexpb.GetIndexStatisticsResponse, error) {
return m.getIndexStatisticsResp, m.err
}
func (m *MockDataCoord) GetIndexInfos(ctx context.Context, req *indexpb.GetIndexInfoRequest) (*indexpb.GetIndexInfoResponse, error) {
return m.getIndexInfosResp, m.err
}
func (m *MockDataCoord) GetIndexState(ctx context.Context, req *indexpb.GetIndexStateRequest) (*indexpb.GetIndexStateResponse, error) {
return m.getIndexStateResp, m.err
}
func (m *MockDataCoord) GetIndexBuildProgress(ctx context.Context, req *indexpb.GetIndexBuildProgressRequest) (*indexpb.GetIndexBuildProgressResponse, error) {
return m.getIndexBuildProgressResp, m.err
}
func (m *MockDataCoord) GetSegmentIndexState(ctx context.Context, req *indexpb.GetSegmentIndexStateRequest) (*indexpb.GetSegmentIndexStateResponse, error) {
return m.getSegmentIndexStateResp, m.err
}
func (m *MockDataCoord) DropIndex(ctx context.Context, req *indexpb.DropIndexRequest) (*commonpb.Status, error) {
return m.dropIndexResp, m.err
}
func Test_NewServer(t *testing.T) {
paramtable.Init()
parameters := []string{"tikv", "etcd"}
for _, v := range parameters {
paramtable.Get().Save(paramtable.Get().MetaStoreCfg.MetaStoreType.Key, v)
ctx := context.Background()
getTiKVClient = func(cfg *paramtable.TiKVConfig) (*txnkv.Client, error) {
return tikv.SetupLocalTxn(), nil
}
defer func() {
getTiKVClient = tikv.GetTiKVClient
}()
mockDataCoord := mocks.NewMockDataCoord(t)
server := NewServer(ctx, nil)
assert.NotNil(t, server)
t.Run("Run", func(t *testing.T) {
server.dataCoord = &MockDataCoord{}
// indexCoord := mocks.NewMockIndexCoord(t)
// indexCoord.EXPECT().Init().Return(nil)
// server.indexCoord = indexCoord
err := server.Run()
assert.NoError(t, err)
})
server.dataCoord = mockDataCoord
t.Run("GetComponentStates", func(t *testing.T) {
server.dataCoord = &MockDataCoord{
states: &milvuspb.ComponentStates{},
}
mockDataCoord.EXPECT().GetComponentStates(mock.Anything, mock.Anything).Return(&milvuspb.ComponentStates{}, nil)
states, err := server.GetComponentStates(ctx, nil)
assert.NoError(t, err)
assert.NotNil(t, states)
})
t.Run("GetTimeTickChannel", func(t *testing.T) {
server.dataCoord = &MockDataCoord{
strResp: &milvuspb.StringResponse{},
}
mockDataCoord.EXPECT().GetTimeTickChannel(mock.Anything, mock.Anything).Return(&milvuspb.StringResponse{}, nil)
resp, err := server.GetTimeTickChannel(ctx, nil)
assert.NoError(t, err)
assert.NotNil(t, resp)
})
t.Run("GetStatisticsChannel", func(t *testing.T) {
server.dataCoord = &MockDataCoord{
strResp: &milvuspb.StringResponse{},
}
mockDataCoord.EXPECT().GetStatisticsChannel(mock.Anything, mock.Anything).Return(&milvuspb.StringResponse{}, nil)
resp, err := server.GetStatisticsChannel(ctx, nil)
assert.NoError(t, err)
assert.NotNil(t, resp)
})
t.Run("GetSegmentInfo", func(t *testing.T) {
server.dataCoord = &MockDataCoord{
infoResp: &datapb.GetSegmentInfoResponse{},
}
mockDataCoord.EXPECT().GetSegmentInfo(mock.Anything, mock.Anything).Return(&datapb.GetSegmentInfoResponse{}, nil)
resp, err := server.GetSegmentInfo(ctx, nil)
assert.NoError(t, err)
assert.NotNil(t, resp)
})
t.Run("Flush", func(t *testing.T) {
server.dataCoord = &MockDataCoord{
flushResp: &datapb.FlushResponse{},
}
mockDataCoord.EXPECT().Flush(mock.Anything, mock.Anything).Return(&datapb.FlushResponse{}, nil)
resp, err := server.Flush(ctx, nil)
assert.NoError(t, err)
assert.NotNil(t, resp)
})
t.Run("AssignSegmentID", func(t *testing.T) {
server.dataCoord = &MockDataCoord{
assignResp: &datapb.AssignSegmentIDResponse{},
}
mockDataCoord.EXPECT().AssignSegmentID(mock.Anything, mock.Anything).Return(&datapb.AssignSegmentIDResponse{}, nil)
resp, err := server.AssignSegmentID(ctx, nil)
assert.NoError(t, err)
assert.NotNil(t, resp)
})
t.Run("GetSegmentStates", func(t *testing.T) {
server.dataCoord = &MockDataCoord{
segStateResp: &datapb.GetSegmentStatesResponse{},
}
mockDataCoord.EXPECT().GetSegmentStates(mock.Anything, mock.Anything).Return(&datapb.GetSegmentStatesResponse{}, nil)
resp, err := server.GetSegmentStates(ctx, nil)
assert.NoError(t, err)
assert.NotNil(t, resp)
})
t.Run("GetInsertBinlogPaths", func(t *testing.T) {
server.dataCoord = &MockDataCoord{
binResp: &datapb.GetInsertBinlogPathsResponse{},
}
mockDataCoord.EXPECT().GetInsertBinlogPaths(mock.Anything, mock.Anything).Return(&datapb.GetInsertBinlogPathsResponse{}, nil)
resp, err := server.GetInsertBinlogPaths(ctx, nil)
assert.NoError(t, err)
assert.NotNil(t, resp)
})
t.Run("GetCollectionStatistics", func(t *testing.T) {
server.dataCoord = &MockDataCoord{
colStatResp: &datapb.GetCollectionStatisticsResponse{},
}
mockDataCoord.EXPECT().GetCollectionStatistics(mock.Anything, mock.Anything).Return(&datapb.GetCollectionStatisticsResponse{}, nil)
resp, err := server.GetCollectionStatistics(ctx, nil)
assert.NoError(t, err)
assert.NotNil(t, resp)
})
t.Run("GetPartitionStatistics", func(t *testing.T) {
server.dataCoord = &MockDataCoord{
partStatResp: &datapb.GetPartitionStatisticsResponse{},
}
mockDataCoord.EXPECT().GetPartitionStatistics(mock.Anything, mock.Anything).Return(&datapb.GetPartitionStatisticsResponse{}, nil)
resp, err := server.GetPartitionStatistics(ctx, nil)
assert.NoError(t, err)
assert.NotNil(t, resp)
})
t.Run("GetSegmentInfoChannel", func(t *testing.T) {
server.dataCoord = &MockDataCoord{
strResp: &milvuspb.StringResponse{},
}
mockDataCoord.EXPECT().GetSegmentInfoChannel(mock.Anything, mock.Anything).Return(&milvuspb.StringResponse{}, nil)
resp, err := server.GetSegmentInfoChannel(ctx, nil)
assert.NoError(t, err)
assert.NotNil(t, resp)
})
t.Run("SaveBinlogPaths", func(t *testing.T) {
server.dataCoord = &MockDataCoord{
status: &commonpb.Status{},
}
mockDataCoord.EXPECT().SaveBinlogPaths(mock.Anything, mock.Anything).Return(merr.Success(), nil)
resp, err := server.SaveBinlogPaths(ctx, nil)
assert.NoError(t, err)
assert.NotNil(t, resp)
})
t.Run("GetRecoveryInfo", func(t *testing.T) {
server.dataCoord = &MockDataCoord{
recoverResp: &datapb.GetRecoveryInfoResponse{},
}
mockDataCoord.EXPECT().GetRecoveryInfo(mock.Anything, mock.Anything).Return(&datapb.GetRecoveryInfoResponse{}, nil)
resp, err := server.GetRecoveryInfo(ctx, nil)
assert.NoError(t, err)
assert.NotNil(t, resp)
})
t.Run("GetFlushedSegments", func(t *testing.T) {
server.dataCoord = &MockDataCoord{
flushSegResp: &datapb.GetFlushedSegmentsResponse{},
}
mockDataCoord.EXPECT().GetFlushedSegments(mock.Anything, mock.Anything).Return(&datapb.GetFlushedSegmentsResponse{}, nil)
resp, err := server.GetFlushedSegments(ctx, nil)
assert.NoError(t, err)
assert.NotNil(t, resp)
})
t.Run("ShowConfigurations", func(t *testing.T) {
server.dataCoord = &MockDataCoord{
configResp: &internalpb.ShowConfigurationsResponse{},
}
mockDataCoord.EXPECT().ShowConfigurations(mock.Anything, mock.Anything).Return(&internalpb.ShowConfigurationsResponse{}, nil)
resp, err := server.ShowConfigurations(ctx, nil)
assert.NoError(t, err)
assert.NotNil(t, resp)
})
t.Run("GetMetrics", func(t *testing.T) {
server.dataCoord = &MockDataCoord{
metricResp: &milvuspb.GetMetricsResponse{},
}
mockDataCoord.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(&milvuspb.GetMetricsResponse{}, nil)
resp, err := server.GetMetrics(ctx, nil)
assert.NoError(t, err)
assert.NotNil(t, resp)
})
t.Run("WatchChannels", func(t *testing.T) {
server.dataCoord = &MockDataCoord{
watchChannelsResp: &datapb.WatchChannelsResponse{},
}
mockDataCoord.EXPECT().WatchChannels(mock.Anything, mock.Anything).Return(&datapb.WatchChannelsResponse{}, nil)
resp, err := server.WatchChannels(ctx, nil)
assert.NoError(t, err)
assert.NotNil(t, resp)
})
t.Run("GetFlushState", func(t *testing.T) {
server.dataCoord = &MockDataCoord{
getFlushStateResp: &milvuspb.GetFlushStateResponse{},
}
mockDataCoord.EXPECT().GetFlushState(mock.Anything, mock.Anything).Return(&milvuspb.GetFlushStateResponse{}, nil)
resp, err := server.GetFlushState(ctx, nil)
assert.NoError(t, err)
assert.NotNil(t, resp)
})
t.Run("GetFlushAllState", func(t *testing.T) {
server.dataCoord = &MockDataCoord{
getFlushAllStateResp: &milvuspb.GetFlushAllStateResponse{},
}
mockDataCoord.EXPECT().GetFlushAllState(mock.Anything, mock.Anything).Return(&milvuspb.GetFlushAllStateResponse{}, nil)
resp, err := server.GetFlushAllState(ctx, nil)
assert.NoError(t, err)
assert.NotNil(t, resp)
})
t.Run("DropVirtualChannel", func(t *testing.T) {
server.dataCoord = &MockDataCoord{
dropVChanResp: &datapb.DropVirtualChannelResponse{},
}
mockDataCoord.EXPECT().DropVirtualChannel(mock.Anything, mock.Anything).Return(&datapb.DropVirtualChannelResponse{}, nil)
resp, err := server.DropVirtualChannel(ctx, nil)
assert.NoError(t, err)
assert.NotNil(t, resp)
})
t.Run("ManualCompaction", func(t *testing.T) {
server.dataCoord = &MockDataCoord{
manualCompactionResp: &milvuspb.ManualCompactionResponse{},
}
mockDataCoord.EXPECT().ManualCompaction(mock.Anything, mock.Anything).Return(&milvuspb.ManualCompactionResponse{}, nil)
resp, err := server.ManualCompaction(ctx, nil)
assert.NoError(t, err)
assert.NotNil(t, resp)
})
t.Run("GetCompactionState", func(t *testing.T) {
server.dataCoord = &MockDataCoord{
compactionStateResp: &milvuspb.GetCompactionStateResponse{},
}
mockDataCoord.EXPECT().GetCompactionState(mock.Anything, mock.Anything).Return(&milvuspb.GetCompactionStateResponse{}, nil)
resp, err := server.GetCompactionState(ctx, nil)
assert.NoError(t, err)
assert.NotNil(t, resp)
})
t.Run("GetCompactionStateWithPlans", func(t *testing.T) {
server.dataCoord = &MockDataCoord{
compactionPlansResp: &milvuspb.GetCompactionPlansResponse{},
}
mockDataCoord.EXPECT().GetCompactionStateWithPlans(mock.Anything, mock.Anything).Return(&milvuspb.GetCompactionPlansResponse{}, nil)
resp, err := server.GetCompactionStateWithPlans(ctx, nil)
assert.NoError(t, err)
assert.NotNil(t, resp)
})
t.Run("set segment state", func(t *testing.T) {
server.dataCoord = &MockDataCoord{
setSegmentStateResp: &datapb.SetSegmentStateResponse{},
}
t.Run("SetSegmentState", func(t *testing.T) {
mockDataCoord.EXPECT().SetSegmentState(mock.Anything, mock.Anything).Return(&datapb.SetSegmentStateResponse{}, nil)
resp, err := server.SetSegmentState(ctx, nil)
assert.NoError(t, err)
assert.NotNil(t, resp)
})
t.Run("import", func(t *testing.T) {
server.dataCoord = &MockDataCoord{
importResp: &datapb.ImportTaskResponse{
Status: &commonpb.Status{},
},
}
mockDataCoord.EXPECT().Import(mock.Anything, mock.Anything).Return(&datapb.ImportTaskResponse{}, nil)
resp, err := server.Import(ctx, nil)
assert.NoError(t, err)
assert.NotNil(t, resp)
})
t.Run("update seg stat", func(t *testing.T) {
server.dataCoord = &MockDataCoord{
updateSegStatResp: merr.Success(),
}
t.Run("UpdateSegmentStatistics", func(t *testing.T) {
mockDataCoord.EXPECT().UpdateSegmentStatistics(mock.Anything, mock.Anything).Return(merr.Success(), nil)
resp, err := server.UpdateSegmentStatistics(ctx, nil)
assert.NoError(t, err)
assert.NotNil(t, resp)
})
t.Run("UpdateChannelCheckpoint", func(t *testing.T) {
server.dataCoord = &MockDataCoord{
updateChanPos: merr.Success(),
}
mockDataCoord.EXPECT().UpdateChannelCheckpoint(mock.Anything, mock.Anything).Return(merr.Success(), nil)
resp, err := server.UpdateChannelCheckpoint(ctx, nil)
assert.NoError(t, err)
assert.NotNil(t, resp)
})
t.Run("save import segment", func(t *testing.T) {
server.dataCoord = &MockDataCoord{
addSegmentResp: merr.Success(),
}
t.Run("SaveImportSegment", func(t *testing.T) {
mockDataCoord.EXPECT().SaveImportSegment(mock.Anything, mock.Anything).Return(merr.Success(), nil)
resp, err := server.SaveImportSegment(ctx, nil)
assert.NoError(t, err)
assert.NotNil(t, resp)
})
t.Run("unset isImporting state", func(t *testing.T) {
server.dataCoord = &MockDataCoord{
unsetIsImportingStateResp: merr.Success(),
}
t.Run("UnsetIsImportingState", func(t *testing.T) {
mockDataCoord.EXPECT().UnsetIsImportingState(mock.Anything, mock.Anything).Return(merr.Success(), nil)
resp, err := server.UnsetIsImportingState(ctx, nil)
assert.NoError(t, err)
assert.NotNil(t, resp)
})
t.Run("mark segments dropped", func(t *testing.T) {
server.dataCoord = &MockDataCoord{
markSegmentsDroppedResp: merr.Success(),
}
t.Run("MarkSegmentsDropped", func(t *testing.T) {
mockDataCoord.EXPECT().MarkSegmentsDropped(mock.Anything, mock.Anything).Return(merr.Success(), nil)
resp, err := server.MarkSegmentsDropped(ctx, nil)
assert.NoError(t, err)
assert.NotNil(t, resp)
})
t.Run("broadcast altered collection", func(t *testing.T) {
server.dataCoord = &MockDataCoord{
broadCastResp: &commonpb.Status{},
}
t.Run("BroadcastAlteredCollection", func(t *testing.T) {
mockDataCoord.EXPECT().BroadcastAlteredCollection(mock.Anything, mock.Anything).Return(merr.Success(), nil)
resp, err := server.BroadcastAlteredCollection(ctx, nil)
assert.NoError(t, err)
assert.NotNil(t, resp)
})
t.Run("CheckHealth", func(t *testing.T) {
server.dataCoord = &MockDataCoord{}
mockDataCoord.EXPECT().CheckHealth(mock.Anything, mock.Anything).Return(&milvuspb.CheckHealthResponse{IsHealthy: true}, nil)
ret, err := server.CheckHealth(ctx, nil)
assert.NoError(t, err)
assert.Equal(t, true, ret.IsHealthy)
})
t.Run("CreateIndex", func(t *testing.T) {
server.dataCoord = &MockDataCoord{
createIndexResp: &commonpb.Status{},
}
mockDataCoord.EXPECT().CreateIndex(mock.Anything, mock.Anything).Return(merr.Success(), nil)
ret, err := server.CreateIndex(ctx, nil)
assert.NoError(t, err)
assert.NotNil(t, ret)
})
t.Run("DescribeIndex", func(t *testing.T) {
server.dataCoord = &MockDataCoord{
describeIndexResp: &indexpb.DescribeIndexResponse{},
}
mockDataCoord.EXPECT().DescribeIndex(mock.Anything, mock.Anything).Return(&indexpb.DescribeIndexResponse{}, nil)
ret, err := server.DescribeIndex(ctx, nil)
assert.NoError(t, err)
assert.NotNil(t, ret)
})
t.Run("GetIndexStatistics", func(t *testing.T) {
server.dataCoord = &MockDataCoord{
getIndexStatisticsResp: &indexpb.GetIndexStatisticsResponse{},
}
mockDataCoord.EXPECT().GetIndexStatistics(mock.Anything, mock.Anything).Return(&indexpb.GetIndexStatisticsResponse{}, nil)
ret, err := server.GetIndexStatistics(ctx, nil)
assert.NoError(t, err)
assert.NotNil(t, ret)
})
t.Run("DropIndex", func(t *testing.T) {
server.dataCoord = &MockDataCoord{
dropIndexResp: &commonpb.Status{},
}
mockDataCoord.EXPECT().DropIndex(mock.Anything, mock.Anything).Return(merr.Success(), nil)
ret, err := server.DropIndex(ctx, nil)
assert.NoError(t, err)
assert.NotNil(t, ret)
})
t.Run("GetIndexState", func(t *testing.T) {
server.dataCoord = &MockDataCoord{
getIndexStateResp: &indexpb.GetIndexStateResponse{},
}
mockDataCoord.EXPECT().GetIndexState(mock.Anything, mock.Anything).Return(&indexpb.GetIndexStateResponse{}, nil)
ret, err := server.GetIndexState(ctx, nil)
assert.NoError(t, err)
assert.NotNil(t, ret)
})
t.Run("GetIndexBuildProgress", func(t *testing.T) {
server.dataCoord = &MockDataCoord{
getIndexBuildProgressResp: &indexpb.GetIndexBuildProgressResponse{},
}
mockDataCoord.EXPECT().GetIndexBuildProgress(mock.Anything, mock.Anything).Return(&indexpb.GetIndexBuildProgressResponse{}, nil)
ret, err := server.GetIndexBuildProgress(ctx, nil)
assert.NoError(t, err)
assert.NotNil(t, ret)
})
t.Run("GetSegmentIndexState", func(t *testing.T) {
server.dataCoord = &MockDataCoord{
getSegmentIndexStateResp: &indexpb.GetSegmentIndexStateResponse{},
}
mockDataCoord.EXPECT().GetSegmentIndexState(mock.Anything, mock.Anything).Return(&indexpb.GetSegmentIndexStateResponse{}, nil)
ret, err := server.GetSegmentIndexState(ctx, nil)
assert.NoError(t, err)
assert.NotNil(t, ret)
})
t.Run("GetIndexInfos", func(t *testing.T) {
server.dataCoord = &MockDataCoord{
getIndexInfosResp: &indexpb.GetIndexInfoResponse{},
}
mockDataCoord.EXPECT().GetIndexInfos(mock.Anything, mock.Anything).Return(&indexpb.GetIndexInfoResponse{}, nil)
ret, err := server.GetIndexInfos(ctx, nil)
assert.NoError(t, err)
assert.NotNil(t, ret)
})
err := server.Stop()
assert.NoError(t, err)
}
}
func Test_Run(t *testing.T) {
paramtable.Init()
t.Run("test run success", func(t *testing.T) {
parameters := []string{"tikv", "etcd"}
for _, v := range parameters {
t.Run(fmt.Sprintf("Run server with %s as metadata storage", v), func(t *testing.T) {
paramtable.Get().Save(paramtable.Get().MetaStoreCfg.MetaStoreType.Key, v)
ctx := context.Background()
getTiKVClient = func(cfg *paramtable.TiKVConfig) (*txnkv.Client, error) {
@ -696,33 +342,97 @@ func Test_Run(t *testing.T) {
server := NewServer(ctx, nil)
assert.NotNil(t, server)
server.dataCoord = &MockDataCoord{
regErr: errors.New("error"),
mockDataCoord := mocks.NewMockDataCoord(t)
server.dataCoord = mockDataCoord
mockDataCoord.EXPECT().SetEtcdClient(mock.Anything)
mockDataCoord.EXPECT().SetAddress(mock.Anything)
mockDataCoord.EXPECT().SetTiKVClient(mock.Anything).Maybe()
mockDataCoord.EXPECT().Init().Return(nil)
mockDataCoord.EXPECT().Start().Return(nil)
mockDataCoord.EXPECT().Register().Return(nil)
err := server.Run()
assert.NoError(t, err)
mockDataCoord.EXPECT().Stop().Return(nil)
err = server.Stop()
assert.NoError(t, err)
}
})
paramtable.Get().Save(paramtable.Get().MetaStoreCfg.MetaStoreType.Key, "etcd")
t.Run("test init error", func(t *testing.T) {
ctx := context.Background()
server := NewServer(ctx, nil)
assert.NotNil(t, server)
mockDataCoord := mocks.NewMockDataCoord(t)
mockDataCoord.EXPECT().SetEtcdClient(mock.Anything)
mockDataCoord.EXPECT().SetAddress(mock.Anything)
mockDataCoord.EXPECT().Init().Return(errors.New("error"))
server.dataCoord = mockDataCoord
err := server.Run()
assert.Error(t, err)
server.dataCoord = &MockDataCoord{
startErr: errors.New("error"),
}
mockDataCoord.EXPECT().Stop().Return(nil)
server.Stop()
})
err = server.Run()
t.Run("test register error", func(t *testing.T) {
ctx := context.Background()
server := NewServer(ctx, nil)
assert.NotNil(t, server)
mockDataCoord := mocks.NewMockDataCoord(t)
mockDataCoord.EXPECT().SetEtcdClient(mock.Anything)
mockDataCoord.EXPECT().SetAddress(mock.Anything)
mockDataCoord.EXPECT().Init().Return(nil)
mockDataCoord.EXPECT().Register().Return(errors.New("error"))
server.dataCoord = mockDataCoord
err := server.Run()
assert.Error(t, err)
server.dataCoord = &MockDataCoord{
initErr: errors.New("error"),
}
mockDataCoord.EXPECT().Stop().Return(nil)
server.Stop()
})
err = server.Run()
t.Run("test start error", func(t *testing.T) {
ctx := context.Background()
server := NewServer(ctx, nil)
assert.NotNil(t, server)
mockDataCoord := mocks.NewMockDataCoord(t)
mockDataCoord.EXPECT().SetEtcdClient(mock.Anything)
mockDataCoord.EXPECT().SetAddress(mock.Anything)
mockDataCoord.EXPECT().Init().Return(nil)
mockDataCoord.EXPECT().Register().Return(nil)
mockDataCoord.EXPECT().Start().Return(errors.New("error"))
server.dataCoord = mockDataCoord
err := server.Run()
assert.Error(t, err)
server.dataCoord = &MockDataCoord{
stopErr: errors.New("error"),
}
mockDataCoord.EXPECT().Stop().Return(nil)
server.Stop()
})
t.Run("test stop error", func(t *testing.T) {
ctx := context.Background()
server := NewServer(ctx, nil)
assert.NotNil(t, server)
mockDataCoord := mocks.NewMockDataCoord(t)
mockDataCoord.EXPECT().SetEtcdClient(mock.Anything)
mockDataCoord.EXPECT().SetAddress(mock.Anything)
mockDataCoord.EXPECT().Init().Return(nil)
mockDataCoord.EXPECT().Register().Return(nil)
mockDataCoord.EXPECT().Start().Return(nil)
server.dataCoord = mockDataCoord
err := server.Run()
assert.NoError(t, err)
mockDataCoord.EXPECT().Stop().Return(errors.New("error"))
err = server.Stop()
assert.Error(t, err)
})
}
}

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

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

9
internal/rootcoord/broker_test.go
View File

@ -227,8 +227,8 @@ 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{
mockDataCoord := mocks.NewMockDataCoordClient(t)
mockDataCoord.EXPECT().GetSegmentIndexState(mock.Anything, mock.Anything).Return(&indexpb.GetSegmentIndexStateResponse{
Status: merr.Success(),
States: []*indexpb.SegmentIndexState{
{
@ -237,8 +237,9 @@ func TestServerBroker_GetSegmentIndexState(t *testing.T) {
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
}