admin/milvus
1
0
Fork 0
mirror of https://gitee.com/milvus-io/milvus.git synced 2026-02-04 11:18:44 +08:00
Code Issues Packages Projects Releases Wiki Activity

Organize mocks from types.go (#25466)

Browse Source 
Signed-off-by: yiwangdr <yiwangdr@gmail.com>
This commit is contained in:
yiwangdr 2023-07-13 19:12:31 -07:00 committed by GitHub
parent fc9a6dc241
commit b9189b9f41
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
32 changed files with 4234 additions and 8729 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
internal/datacoord/compaction_test.go
View File

@ -199,7 +199,7 @@ func Test_compactionPlanHandler_execCompactionPlan(t *testing.T) {
func Test_compactionPlanHandler_execWithParallels(t *testing.T) {
mockDataNode := &mocks.DataNode{}
mockDataNode := &mocks.MockDataNode{}
paramtable.Get().Save(Params.DataCoordCfg.CompactionCheckIntervalInSeconds.Key, "1")
defer paramtable.Get().Reset(Params.DataCoordCfg.CompactionCheckIntervalInSeconds.Key)
c := &compactionPlanHandler{
@ -286,7 +286,7 @@ func getDeltaLogPath(rootPath string, segmentID typeutil.UniqueID) string {
}
func TestCompactionPlanHandler_handleMergeCompactionResult(t *testing.T) {
mockDataNode := &mocks.DataNode{}
mockDataNode := &mocks.MockDataNode{}
call := mockDataNode.EXPECT().SyncSegments(mock.Anything, mock.Anything).Run(func(ctx context.Context, req *datapb.SyncSegmentsRequest) {}).Return(&commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}, nil)
dataNodeID := UniqueID(111)
@ -441,7 +441,7 @@ func TestCompactionPlanHandler_completeCompaction(t *testing.T) {
})
t.Run("test complete merge compaction task", func(t *testing.T) {
mockDataNode := &mocks.DataNode{}
mockDataNode := &mocks.MockDataNode{}
mockDataNode.EXPECT().SyncSegments(mock.Anything, mock.Anything).Run(func(ctx context.Context, req *datapb.SyncSegmentsRequest) {}).Return(&commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}, nil)
dataNodeID := UniqueID(111)
@ -533,7 +533,7 @@ func TestCompactionPlanHandler_completeCompaction(t *testing.T) {
})
t.Run("test empty result merge compaction task", func(t *testing.T) {
mockDataNode := &mocks.DataNode{}
mockDataNode := &mocks.MockDataNode{}
mockDataNode.EXPECT().SyncSegments(mock.Anything, mock.Anything).Run(func(ctx context.Context, req *datapb.SyncSegmentsRequest) {}).Return(&commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}, nil)
dataNodeID := UniqueID(111)

4
internal/datanode/timetick_sender_test.go
View File

@ -26,8 +26,8 @@ import (
"go.uber.org/atomic"
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus/internal/mocks"
"github.com/milvus-io/milvus/internal/proto/datapb"
"github.com/milvus-io/milvus/internal/types"
)
func TestTimetickManagerNormal(t *testing.T) {
@ -166,7 +166,7 @@ func TestTimetickManagerSendNotSuccess(t *testing.T) {
func TestTimetickManagerSendReport(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
mockDataCoord := types.NewMockDataCoord(t)
mockDataCoord := mocks.NewMockDataCoord(t)
tsInMill := time.Now().UnixMilli()
validTs := atomic.NewBool(false)

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

@ -42,6 +42,7 @@ import (
"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/proto/proxypb"
@ -1526,7 +1527,7 @@ func getServer(t *testing.T) *Server {
server.rootCoordClient = &MockRootCoord{}
server.dataCoordClient = &MockDataCoord{}
mockQC := &types.MockQueryCoord{}
mockQC := &mocks.MockQueryCoord{}
server.queryCoordClient = mockQC
mockQC.EXPECT().Init().Return(nil)
mockQC.EXPECT().GetComponentStates(mock.Anything).Return(&milvuspb.ComponentStates{

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

@ -23,6 +23,7 @@ import (
"github.com/cockroachdb/errors"
"github.com/milvus-io/milvus/internal/mocks"
"github.com/milvus-io/milvus/internal/types"
"github.com/milvus-io/milvus/pkg/util/paramtable"
@ -349,8 +350,8 @@ func TestServer_Run2(t *testing.T) {
assert.NoError(t, err)
}
func getQueryCoord() *types.MockQueryCoord {
mqc := &types.MockQueryCoord{}
func getQueryCoord() *mocks.MockQueryCoord {
mqc := &mocks.MockQueryCoord{}
mqc.EXPECT().Init().Return(nil)
mqc.EXPECT().SetEtcdClient(mock.Anything)
mqc.EXPECT().SetAddress(mock.Anything)

5
internal/distributed/querynode/service_test.go
View File

@ -23,6 +23,7 @@ import (
"github.com/cockroachdb/errors"
"github.com/milvus-io/milvus/internal/mocks"
"github.com/milvus-io/milvus/internal/types"
"github.com/milvus-io/milvus/pkg/util/paramtable"
@ -84,7 +85,7 @@ func Test_NewServer(t *testing.T) {
assert.NoError(t, err)
assert.NotNil(t, server)
mockQN := types.NewMockQueryNode(t)
mockQN := mocks.NewMockQueryNode(t)
mockQN.EXPECT().Start().Return(nil).Maybe()
mockQN.EXPECT().Stop().Return(nil).Maybe()
mockQN.EXPECT().Register().Return(nil).Maybe()
@ -260,7 +261,7 @@ func Test_Run(t *testing.T) {
assert.NoError(t, err)
assert.NotNil(t, server)
mockQN := types.NewMockQueryNode(t)
mockQN := mocks.NewMockQueryNode(t)
mockQN.EXPECT().Start().Return(errors.New("Failed")).Maybe()
mockQN.EXPECT().Stop().Return(errors.New("Failed")).Maybe()
mockQN.EXPECT().Register().Return(errors.New("Failed")).Maybe()

1491
internal/mocks/mock_datacoord.go
View File

File diff suppressed because it is too large Load Diff

655
internal/mocks/mock_datanode.go
View File

File diff suppressed because it is too large Load Diff

2267
internal/mocks/mock_proxy.go
View File

File diff suppressed because it is too large Load Diff

427
internal/types/mock_querycoord.go → internal/mocks/mock_querycoord.go
View File

File diff suppressed because it is too large Load Diff

359
internal/types/mock_querynode.go → internal/mocks/mock_querynode.go
View File

@ -1,6 +1,6 @@
// Code generated by mockery v2.16.0. DO NOT EDIT.
// Code generated by mockery v2.30.1. DO NOT EDIT.
package types
package mocks
import (
context "context"
@ -35,6 +35,10 @@ func (_m *MockQueryNode) Delete(_a0 context.Context, _a1 *querypb.DeleteRequest)
ret := _m.Called(_a0, _a1)
var r0 *commonpb.Status
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, *querypb.DeleteRequest) (*commonpb.Status, error)); ok {
return rf(_a0, _a1)
}
if rf, ok := ret.Get(0).(func(context.Context, *querypb.DeleteRequest) *commonpb.Status); ok {
r0 = rf(_a0, _a1)
} else {
@ -43,7 +47,6 @@ func (_m *MockQueryNode) Delete(_a0 context.Context, _a1 *querypb.DeleteRequest)
}
}
var r1 error
if rf, ok := ret.Get(1).(func(context.Context, *querypb.DeleteRequest) error); ok {
r1 = rf(_a0, _a1)
} else {
@ -59,8 +62,8 @@ type MockQueryNode_Delete_Call struct {
}
// Delete is a helper method to define mock.On call
// - _a0 context.Context
// - _a1 *querypb.DeleteRequest
// - _a0 context.Context
// - _a1 *querypb.DeleteRequest
func (_e *MockQueryNode_Expecter) Delete(_a0 interface{}, _a1 interface{}) *MockQueryNode_Delete_Call {
return &MockQueryNode_Delete_Call{Call: _e.mock.On("Delete", _a0, _a1)}
}
@ -77,6 +80,11 @@ func (_c *MockQueryNode_Delete_Call) Return(_a0 *commonpb.Status, _a1 error) *Mo
return _c
}
func (_c *MockQueryNode_Delete_Call) RunAndReturn(run func(context.Context, *querypb.DeleteRequest) (*commonpb.Status, error)) *MockQueryNode_Delete_Call {
_c.Call.Return(run)
return _c
}
// GetAddress provides a mock function with given fields:
func (_m *MockQueryNode) GetAddress() string {
ret := _m.Called()
@ -113,11 +121,20 @@ func (_c *MockQueryNode_GetAddress_Call) Return(_a0 string) *MockQueryNode_GetAd
return _c
}
func (_c *MockQueryNode_GetAddress_Call) RunAndReturn(run func() string) *MockQueryNode_GetAddress_Call {
_c.Call.Return(run)
return _c
}
// GetComponentStates provides a mock function with given fields: ctx
func (_m *MockQueryNode) GetComponentStates(ctx context.Context) (*milvuspb.ComponentStates, error) {
ret := _m.Called(ctx)
var r0 *milvuspb.ComponentStates
var r1 error
if rf, ok := ret.Get(0).(func(context.Context) (*milvuspb.ComponentStates, error)); ok {
return rf(ctx)
}
if rf, ok := ret.Get(0).(func(context.Context) *milvuspb.ComponentStates); ok {
r0 = rf(ctx)
} else {
@ -126,7 +143,6 @@ func (_m *MockQueryNode) GetComponentStates(ctx context.Context) (*milvuspb.Comp
}
}
var r1 error
if rf, ok := ret.Get(1).(func(context.Context) error); ok {
r1 = rf(ctx)
} else {
@ -142,7 +158,7 @@ type MockQueryNode_GetComponentStates_Call struct {
}
// GetComponentStates is a helper method to define mock.On call
// - ctx context.Context
// - ctx context.Context
func (_e *MockQueryNode_Expecter) GetComponentStates(ctx interface{}) *MockQueryNode_GetComponentStates_Call {
return &MockQueryNode_GetComponentStates_Call{Call: _e.mock.On("GetComponentStates", ctx)}
}
@ -159,11 +175,20 @@ func (_c *MockQueryNode_GetComponentStates_Call) Return(_a0 *milvuspb.ComponentS
return _c
}
func (_c *MockQueryNode_GetComponentStates_Call) RunAndReturn(run func(context.Context) (*milvuspb.ComponentStates, error)) *MockQueryNode_GetComponentStates_Call {
_c.Call.Return(run)
return _c
}
// GetDataDistribution provides a mock function with given fields: _a0, _a1
func (_m *MockQueryNode) GetDataDistribution(_a0 context.Context, _a1 *querypb.GetDataDistributionRequest) (*querypb.GetDataDistributionResponse, error) {
ret := _m.Called(_a0, _a1)
var r0 *querypb.GetDataDistributionResponse
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, *querypb.GetDataDistributionRequest) (*querypb.GetDataDistributionResponse, error)); ok {
return rf(_a0, _a1)
}
if rf, ok := ret.Get(0).(func(context.Context, *querypb.GetDataDistributionRequest) *querypb.GetDataDistributionResponse); ok {
r0 = rf(_a0, _a1)
} else {
@ -172,7 +197,6 @@ func (_m *MockQueryNode) GetDataDistribution(_a0 context.Context, _a1 *querypb.G
}
}
var r1 error
if rf, ok := ret.Get(1).(func(context.Context, *querypb.GetDataDistributionRequest) error); ok {
r1 = rf(_a0, _a1)
} else {
@ -188,8 +212,8 @@ type MockQueryNode_GetDataDistribution_Call struct {
}
// GetDataDistribution is a helper method to define mock.On call
// - _a0 context.Context
// - _a1 *querypb.GetDataDistributionRequest
// - _a0 context.Context
// - _a1 *querypb.GetDataDistributionRequest
func (_e *MockQueryNode_Expecter) GetDataDistribution(_a0 interface{}, _a1 interface{}) *MockQueryNode_GetDataDistribution_Call {
return &MockQueryNode_GetDataDistribution_Call{Call: _e.mock.On("GetDataDistribution", _a0, _a1)}
}
@ -206,11 +230,20 @@ func (_c *MockQueryNode_GetDataDistribution_Call) Return(_a0 *querypb.GetDataDis
return _c
}
func (_c *MockQueryNode_GetDataDistribution_Call) RunAndReturn(run func(context.Context, *querypb.GetDataDistributionRequest) (*querypb.GetDataDistributionResponse, error)) *MockQueryNode_GetDataDistribution_Call {
_c.Call.Return(run)
return _c
}
// GetMetrics provides a mock function with given fields: ctx, req
func (_m *MockQueryNode) GetMetrics(ctx context.Context, req *milvuspb.GetMetricsRequest) (*milvuspb.GetMetricsResponse, error) {
ret := _m.Called(ctx, req)
var r0 *milvuspb.GetMetricsResponse
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, *milvuspb.GetMetricsRequest) (*milvuspb.GetMetricsResponse, error)); ok {
return rf(ctx, req)
}
if rf, ok := ret.Get(0).(func(context.Context, *milvuspb.GetMetricsRequest) *milvuspb.GetMetricsResponse); ok {
r0 = rf(ctx, req)
} else {
@ -219,7 +252,6 @@ func (_m *MockQueryNode) GetMetrics(ctx context.Context, req *milvuspb.GetMetric
}
}
var r1 error
if rf, ok := ret.Get(1).(func(context.Context, *milvuspb.GetMetricsRequest) error); ok {
r1 = rf(ctx, req)
} else {
@ -235,8 +267,8 @@ type MockQueryNode_GetMetrics_Call struct {
}
// GetMetrics is a helper method to define mock.On call
// - ctx context.Context
// - req *milvuspb.GetMetricsRequest
// - ctx context.Context
// - req *milvuspb.GetMetricsRequest
func (_e *MockQueryNode_Expecter) GetMetrics(ctx interface{}, req interface{}) *MockQueryNode_GetMetrics_Call {
return &MockQueryNode_GetMetrics_Call{Call: _e.mock.On("GetMetrics", ctx, req)}
}
@ -253,11 +285,20 @@ func (_c *MockQueryNode_GetMetrics_Call) Return(_a0 *milvuspb.GetMetricsResponse
return _c
}
func (_c *MockQueryNode_GetMetrics_Call) RunAndReturn(run func(context.Context, *milvuspb.GetMetricsRequest) (*milvuspb.GetMetricsResponse, error)) *MockQueryNode_GetMetrics_Call {
_c.Call.Return(run)
return _c
}
// GetSegmentInfo provides a mock function with given fields: ctx, req
func (_m *MockQueryNode) GetSegmentInfo(ctx context.Context, req *querypb.GetSegmentInfoRequest) (*querypb.GetSegmentInfoResponse, error) {
ret := _m.Called(ctx, req)
var r0 *querypb.GetSegmentInfoResponse
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, *querypb.GetSegmentInfoRequest) (*querypb.GetSegmentInfoResponse, error)); ok {
return rf(ctx, req)
}
if rf, ok := ret.Get(0).(func(context.Context, *querypb.GetSegmentInfoRequest) *querypb.GetSegmentInfoResponse); ok {
r0 = rf(ctx, req)
} else {
@ -266,7 +307,6 @@ func (_m *MockQueryNode) GetSegmentInfo(ctx context.Context, req *querypb.GetSeg
}
}
var r1 error
if rf, ok := ret.Get(1).(func(context.Context, *querypb.GetSegmentInfoRequest) error); ok {
r1 = rf(ctx, req)
} else {
@ -282,8 +322,8 @@ type MockQueryNode_GetSegmentInfo_Call struct {
}
// GetSegmentInfo is a helper method to define mock.On call
// - ctx context.Context
// - req *querypb.GetSegmentInfoRequest
// - ctx context.Context
// - req *querypb.GetSegmentInfoRequest
func (_e *MockQueryNode_Expecter) GetSegmentInfo(ctx interface{}, req interface{}) *MockQueryNode_GetSegmentInfo_Call {
return &MockQueryNode_GetSegmentInfo_Call{Call: _e.mock.On("GetSegmentInfo", ctx, req)}
}
@ -300,11 +340,20 @@ func (_c *MockQueryNode_GetSegmentInfo_Call) Return(_a0 *querypb.GetSegmentInfoR
return _c
}
func (_c *MockQueryNode_GetSegmentInfo_Call) RunAndReturn(run func(context.Context, *querypb.GetSegmentInfoRequest) (*querypb.GetSegmentInfoResponse, error)) *MockQueryNode_GetSegmentInfo_Call {
_c.Call.Return(run)
return _c
}
// GetStatistics provides a mock function with given fields: ctx, req
func (_m *MockQueryNode) GetStatistics(ctx context.Context, req *querypb.GetStatisticsRequest) (*internalpb.GetStatisticsResponse, error) {
ret := _m.Called(ctx, req)
var r0 *internalpb.GetStatisticsResponse
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, *querypb.GetStatisticsRequest) (*internalpb.GetStatisticsResponse, error)); ok {
return rf(ctx, req)
}
if rf, ok := ret.Get(0).(func(context.Context, *querypb.GetStatisticsRequest) *internalpb.GetStatisticsResponse); ok {
r0 = rf(ctx, req)
} else {
@ -313,7 +362,6 @@ func (_m *MockQueryNode) GetStatistics(ctx context.Context, req *querypb.GetStat
}
}
var r1 error
if rf, ok := ret.Get(1).(func(context.Context, *querypb.GetStatisticsRequest) error); ok {
r1 = rf(ctx, req)
} else {
@ -329,8 +377,8 @@ type MockQueryNode_GetStatistics_Call struct {
}
// GetStatistics is a helper method to define mock.On call
// - ctx context.Context
// - req *querypb.GetStatisticsRequest
// - ctx context.Context
// - req *querypb.GetStatisticsRequest
func (_e *MockQueryNode_Expecter) GetStatistics(ctx interface{}, req interface{}) *MockQueryNode_GetStatistics_Call {
return &MockQueryNode_GetStatistics_Call{Call: _e.mock.On("GetStatistics", ctx, req)}
}
@ -347,11 +395,20 @@ func (_c *MockQueryNode_GetStatistics_Call) Return(_a0 *internalpb.GetStatistics
return _c
}
func (_c *MockQueryNode_GetStatistics_Call) RunAndReturn(run func(context.Context, *querypb.GetStatisticsRequest) (*internalpb.GetStatisticsResponse, error)) *MockQueryNode_GetStatistics_Call {
_c.Call.Return(run)
return _c
}
// GetStatisticsChannel provides a mock function with given fields: ctx
func (_m *MockQueryNode) GetStatisticsChannel(ctx context.Context) (*milvuspb.StringResponse, error) {
ret := _m.Called(ctx)
var r0 *milvuspb.StringResponse
var r1 error
if rf, ok := ret.Get(0).(func(context.Context) (*milvuspb.StringResponse, error)); ok {
return rf(ctx)
}
if rf, ok := ret.Get(0).(func(context.Context) *milvuspb.StringResponse); ok {
r0 = rf(ctx)
} else {
@ -360,7 +417,6 @@ func (_m *MockQueryNode) GetStatisticsChannel(ctx context.Context) (*milvuspb.St
}
}
var r1 error
if rf, ok := ret.Get(1).(func(context.Context) error); ok {
r1 = rf(ctx)
} else {
@ -376,7 +432,7 @@ type MockQueryNode_GetStatisticsChannel_Call struct {
}
// GetStatisticsChannel is a helper method to define mock.On call
// - ctx context.Context
// - ctx context.Context
func (_e *MockQueryNode_Expecter) GetStatisticsChannel(ctx interface{}) *MockQueryNode_GetStatisticsChannel_Call {
return &MockQueryNode_GetStatisticsChannel_Call{Call: _e.mock.On("GetStatisticsChannel", ctx)}
}
@ -393,11 +449,20 @@ func (_c *MockQueryNode_GetStatisticsChannel_Call) Return(_a0 *milvuspb.StringRe
return _c
}
func (_c *MockQueryNode_GetStatisticsChannel_Call) RunAndReturn(run func(context.Context) (*milvuspb.StringResponse, error)) *MockQueryNode_GetStatisticsChannel_Call {
_c.Call.Return(run)
return _c
}
// GetTimeTickChannel provides a mock function with given fields: ctx
func (_m *MockQueryNode) GetTimeTickChannel(ctx context.Context) (*milvuspb.StringResponse, error) {
ret := _m.Called(ctx)
var r0 *milvuspb.StringResponse
var r1 error
if rf, ok := ret.Get(0).(func(context.Context) (*milvuspb.StringResponse, error)); ok {
return rf(ctx)
}
if rf, ok := ret.Get(0).(func(context.Context) *milvuspb.StringResponse); ok {
r0 = rf(ctx)
} else {
@ -406,7 +471,6 @@ func (_m *MockQueryNode) GetTimeTickChannel(ctx context.Context) (*milvuspb.Stri
}
}
var r1 error
if rf, ok := ret.Get(1).(func(context.Context) error); ok {
r1 = rf(ctx)
} else {
@ -422,7 +486,7 @@ type MockQueryNode_GetTimeTickChannel_Call struct {
}
// GetTimeTickChannel is a helper method to define mock.On call
// - ctx context.Context
// - ctx context.Context
func (_e *MockQueryNode_Expecter) GetTimeTickChannel(ctx interface{}) *MockQueryNode_GetTimeTickChannel_Call {
return &MockQueryNode_GetTimeTickChannel_Call{Call: _e.mock.On("GetTimeTickChannel", ctx)}
}
@ -439,6 +503,11 @@ func (_c *MockQueryNode_GetTimeTickChannel_Call) Return(_a0 *milvuspb.StringResp
return _c
}
func (_c *MockQueryNode_GetTimeTickChannel_Call) RunAndReturn(run func(context.Context) (*milvuspb.StringResponse, error)) *MockQueryNode_GetTimeTickChannel_Call {
_c.Call.Return(run)
return _c
}
// Init provides a mock function with given fields:
func (_m *MockQueryNode) Init() error {
ret := _m.Called()
@ -475,11 +544,20 @@ func (_c *MockQueryNode_Init_Call) Return(_a0 error) *MockQueryNode_Init_Call {
return _c
}
func (_c *MockQueryNode_Init_Call) RunAndReturn(run func() error) *MockQueryNode_Init_Call {
_c.Call.Return(run)
return _c
}
// LoadPartitions provides a mock function with given fields: ctx, req
func (_m *MockQueryNode) LoadPartitions(ctx context.Context, req *querypb.LoadPartitionsRequest) (*commonpb.Status, error) {
ret := _m.Called(ctx, req)
var r0 *commonpb.Status
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, *querypb.LoadPartitionsRequest) (*commonpb.Status, error)); ok {
return rf(ctx, req)
}
if rf, ok := ret.Get(0).(func(context.Context, *querypb.LoadPartitionsRequest) *commonpb.Status); ok {
r0 = rf(ctx, req)
} else {
@ -488,7 +566,6 @@ func (_m *MockQueryNode) LoadPartitions(ctx context.Context, req *querypb.LoadPa
}
}
var r1 error
if rf, ok := ret.Get(1).(func(context.Context, *querypb.LoadPartitionsRequest) error); ok {
r1 = rf(ctx, req)
} else {
@ -504,8 +581,8 @@ type MockQueryNode_LoadPartitions_Call struct {
}
// LoadPartitions is a helper method to define mock.On call
// - ctx context.Context
// - req *querypb.LoadPartitionsRequest
// - ctx context.Context
// - req *querypb.LoadPartitionsRequest
func (_e *MockQueryNode_Expecter) LoadPartitions(ctx interface{}, req interface{}) *MockQueryNode_LoadPartitions_Call {
return &MockQueryNode_LoadPartitions_Call{Call: _e.mock.On("LoadPartitions", ctx, req)}
}
@ -522,11 +599,20 @@ func (_c *MockQueryNode_LoadPartitions_Call) Return(_a0 *commonpb.Status, _a1 er
return _c
}
func (_c *MockQueryNode_LoadPartitions_Call) RunAndReturn(run func(context.Context, *querypb.LoadPartitionsRequest) (*commonpb.Status, error)) *MockQueryNode_LoadPartitions_Call {
_c.Call.Return(run)
return _c
}
// LoadSegments provides a mock function with given fields: ctx, req
func (_m *MockQueryNode) LoadSegments(ctx context.Context, req *querypb.LoadSegmentsRequest) (*commonpb.Status, error) {
ret := _m.Called(ctx, req)
var r0 *commonpb.Status
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, *querypb.LoadSegmentsRequest) (*commonpb.Status, error)); ok {
return rf(ctx, req)
}
if rf, ok := ret.Get(0).(func(context.Context, *querypb.LoadSegmentsRequest) *commonpb.Status); ok {
r0 = rf(ctx, req)
} else {
@ -535,7 +621,6 @@ func (_m *MockQueryNode) LoadSegments(ctx context.Context, req *querypb.LoadSegm
}
}
var r1 error
if rf, ok := ret.Get(1).(func(context.Context, *querypb.LoadSegmentsRequest) error); ok {
r1 = rf(ctx, req)
} else {
@ -551,8 +636,8 @@ type MockQueryNode_LoadSegments_Call struct {
}
// LoadSegments is a helper method to define mock.On call
// - ctx context.Context
// - req *querypb.LoadSegmentsRequest
// - ctx context.Context
// - req *querypb.LoadSegmentsRequest
func (_e *MockQueryNode_Expecter) LoadSegments(ctx interface{}, req interface{}) *MockQueryNode_LoadSegments_Call {
return &MockQueryNode_LoadSegments_Call{Call: _e.mock.On("LoadSegments", ctx, req)}
}
@ -569,11 +654,20 @@ func (_c *MockQueryNode_LoadSegments_Call) Return(_a0 *commonpb.Status, _a1 erro
return _c
}
func (_c *MockQueryNode_LoadSegments_Call) RunAndReturn(run func(context.Context, *querypb.LoadSegmentsRequest) (*commonpb.Status, error)) *MockQueryNode_LoadSegments_Call {
_c.Call.Return(run)
return _c
}
// Query provides a mock function with given fields: ctx, req
func (_m *MockQueryNode) Query(ctx context.Context, req *querypb.QueryRequest) (*internalpb.RetrieveResults, error) {
ret := _m.Called(ctx, req)
var r0 *internalpb.RetrieveResults
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, *querypb.QueryRequest) (*internalpb.RetrieveResults, error)); ok {
return rf(ctx, req)
}
if rf, ok := ret.Get(0).(func(context.Context, *querypb.QueryRequest) *internalpb.RetrieveResults); ok {
r0 = rf(ctx, req)
} else {
@ -582,7 +676,6 @@ func (_m *MockQueryNode) Query(ctx context.Context, req *querypb.QueryRequest) (
}
}
var r1 error
if rf, ok := ret.Get(1).(func(context.Context, *querypb.QueryRequest) error); ok {
r1 = rf(ctx, req)
} else {
@ -598,8 +691,8 @@ type MockQueryNode_Query_Call struct {
}
// Query is a helper method to define mock.On call
// - ctx context.Context
// - req *querypb.QueryRequest
// - ctx context.Context
// - req *querypb.QueryRequest
func (_e *MockQueryNode_Expecter) Query(ctx interface{}, req interface{}) *MockQueryNode_Query_Call {
return &MockQueryNode_Query_Call{Call: _e.mock.On("Query", ctx, req)}
}
@ -616,11 +709,20 @@ func (_c *MockQueryNode_Query_Call) Return(_a0 *internalpb.RetrieveResults, _a1
return _c
}
func (_c *MockQueryNode_Query_Call) RunAndReturn(run func(context.Context, *querypb.QueryRequest) (*internalpb.RetrieveResults, error)) *MockQueryNode_Query_Call {
_c.Call.Return(run)
return _c
}
// QuerySegments provides a mock function with given fields: ctx, req
func (_m *MockQueryNode) QuerySegments(ctx context.Context, req *querypb.QueryRequest) (*internalpb.RetrieveResults, error) {
ret := _m.Called(ctx, req)
var r0 *internalpb.RetrieveResults
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, *querypb.QueryRequest) (*internalpb.RetrieveResults, error)); ok {
return rf(ctx, req)
}
if rf, ok := ret.Get(0).(func(context.Context, *querypb.QueryRequest) *internalpb.RetrieveResults); ok {
r0 = rf(ctx, req)
} else {
@ -629,7 +731,6 @@ func (_m *MockQueryNode) QuerySegments(ctx context.Context, req *querypb.QueryRe
}
}
var r1 error
if rf, ok := ret.Get(1).(func(context.Context, *querypb.QueryRequest) error); ok {
r1 = rf(ctx, req)
} else {
@ -645,8 +746,8 @@ type MockQueryNode_QuerySegments_Call struct {
}
// QuerySegments is a helper method to define mock.On call
// - ctx context.Context
// - req *querypb.QueryRequest
// - ctx context.Context
// - req *querypb.QueryRequest
func (_e *MockQueryNode_Expecter) QuerySegments(ctx interface{}, req interface{}) *MockQueryNode_QuerySegments_Call {
return &MockQueryNode_QuerySegments_Call{Call: _e.mock.On("QuerySegments", ctx, req)}
}
@ -663,6 +764,11 @@ func (_c *MockQueryNode_QuerySegments_Call) Return(_a0 *internalpb.RetrieveResul
return _c
}
func (_c *MockQueryNode_QuerySegments_Call) RunAndReturn(run func(context.Context, *querypb.QueryRequest) (*internalpb.RetrieveResults, error)) *MockQueryNode_QuerySegments_Call {
_c.Call.Return(run)
return _c
}
// Register provides a mock function with given fields:
func (_m *MockQueryNode) Register() error {
ret := _m.Called()
@ -699,11 +805,20 @@ func (_c *MockQueryNode_Register_Call) Return(_a0 error) *MockQueryNode_Register
return _c
}
func (_c *MockQueryNode_Register_Call) RunAndReturn(run func() error) *MockQueryNode_Register_Call {
_c.Call.Return(run)
return _c
}
// ReleaseCollection provides a mock function with given fields: ctx, req
func (_m *MockQueryNode) ReleaseCollection(ctx context.Context, req *querypb.ReleaseCollectionRequest) (*commonpb.Status, error) {
ret := _m.Called(ctx, req)
var r0 *commonpb.Status
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, *querypb.ReleaseCollectionRequest) (*commonpb.Status, error)); ok {
return rf(ctx, req)
}
if rf, ok := ret.Get(0).(func(context.Context, *querypb.ReleaseCollectionRequest) *commonpb.Status); ok {
r0 = rf(ctx, req)
} else {
@ -712,7 +827,6 @@ func (_m *MockQueryNode) ReleaseCollection(ctx context.Context, req *querypb.Rel
}
}
var r1 error
if rf, ok := ret.Get(1).(func(context.Context, *querypb.ReleaseCollectionRequest) error); ok {
r1 = rf(ctx, req)
} else {
@ -728,8 +842,8 @@ type MockQueryNode_ReleaseCollection_Call struct {
}
// ReleaseCollection is a helper method to define mock.On call
// - ctx context.Context
// - req *querypb.ReleaseCollectionRequest
// - ctx context.Context
// - req *querypb.ReleaseCollectionRequest
func (_e *MockQueryNode_Expecter) ReleaseCollection(ctx interface{}, req interface{}) *MockQueryNode_ReleaseCollection_Call {
return &MockQueryNode_ReleaseCollection_Call{Call: _e.mock.On("ReleaseCollection", ctx, req)}
}
@ -746,11 +860,20 @@ func (_c *MockQueryNode_ReleaseCollection_Call) Return(_a0 *commonpb.Status, _a1
return _c
}
func (_c *MockQueryNode_ReleaseCollection_Call) RunAndReturn(run func(context.Context, *querypb.ReleaseCollectionRequest) (*commonpb.Status, error)) *MockQueryNode_ReleaseCollection_Call {
_c.Call.Return(run)
return _c
}
// ReleasePartitions provides a mock function with given fields: ctx, req
func (_m *MockQueryNode) ReleasePartitions(ctx context.Context, req *querypb.ReleasePartitionsRequest) (*commonpb.Status, error) {
ret := _m.Called(ctx, req)
var r0 *commonpb.Status
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, *querypb.ReleasePartitionsRequest) (*commonpb.Status, error)); ok {
return rf(ctx, req)
}
if rf, ok := ret.Get(0).(func(context.Context, *querypb.ReleasePartitionsRequest) *commonpb.Status); ok {
r0 = rf(ctx, req)
} else {
@ -759,7 +882,6 @@ func (_m *MockQueryNode) ReleasePartitions(ctx context.Context, req *querypb.Rel
}
}
var r1 error
if rf, ok := ret.Get(1).(func(context.Context, *querypb.ReleasePartitionsRequest) error); ok {
r1 = rf(ctx, req)
} else {
@ -775,8 +897,8 @@ type MockQueryNode_ReleasePartitions_Call struct {
}
// ReleasePartitions is a helper method to define mock.On call
// - ctx context.Context
// - req *querypb.ReleasePartitionsRequest
// - ctx context.Context
// - req *querypb.ReleasePartitionsRequest
func (_e *MockQueryNode_Expecter) ReleasePartitions(ctx interface{}, req interface{}) *MockQueryNode_ReleasePartitions_Call {
return &MockQueryNode_ReleasePartitions_Call{Call: _e.mock.On("ReleasePartitions", ctx, req)}
}
@ -793,11 +915,20 @@ func (_c *MockQueryNode_ReleasePartitions_Call) Return(_a0 *commonpb.Status, _a1
return _c
}
func (_c *MockQueryNode_ReleasePartitions_Call) RunAndReturn(run func(context.Context, *querypb.ReleasePartitionsRequest) (*commonpb.Status, error)) *MockQueryNode_ReleasePartitions_Call {
_c.Call.Return(run)
return _c
}
// ReleaseSegments provides a mock function with given fields: ctx, req
func (_m *MockQueryNode) ReleaseSegments(ctx context.Context, req *querypb.ReleaseSegmentsRequest) (*commonpb.Status, error) {
ret := _m.Called(ctx, req)
var r0 *commonpb.Status
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, *querypb.ReleaseSegmentsRequest) (*commonpb.Status, error)); ok {
return rf(ctx, req)
}
if rf, ok := ret.Get(0).(func(context.Context, *querypb.ReleaseSegmentsRequest) *commonpb.Status); ok {
r0 = rf(ctx, req)
} else {
@ -806,7 +937,6 @@ func (_m *MockQueryNode) ReleaseSegments(ctx context.Context, req *querypb.Relea
}
}
var r1 error
if rf, ok := ret.Get(1).(func(context.Context, *querypb.ReleaseSegmentsRequest) error); ok {
r1 = rf(ctx, req)
} else {
@ -822,8 +952,8 @@ type MockQueryNode_ReleaseSegments_Call struct {
}
// ReleaseSegments is a helper method to define mock.On call
// - ctx context.Context
// - req *querypb.ReleaseSegmentsRequest
// - ctx context.Context
// - req *querypb.ReleaseSegmentsRequest
func (_e *MockQueryNode_Expecter) ReleaseSegments(ctx interface{}, req interface{}) *MockQueryNode_ReleaseSegments_Call {
return &MockQueryNode_ReleaseSegments_Call{Call: _e.mock.On("ReleaseSegments", ctx, req)}
}
@ -840,11 +970,20 @@ func (_c *MockQueryNode_ReleaseSegments_Call) Return(_a0 *commonpb.Status, _a1 e
return _c
}
func (_c *MockQueryNode_ReleaseSegments_Call) RunAndReturn(run func(context.Context, *querypb.ReleaseSegmentsRequest) (*commonpb.Status, error)) *MockQueryNode_ReleaseSegments_Call {
_c.Call.Return(run)
return _c
}
// Search provides a mock function with given fields: ctx, req
func (_m *MockQueryNode) Search(ctx context.Context, req *querypb.SearchRequest) (*internalpb.SearchResults, error) {
ret := _m.Called(ctx, req)
var r0 *internalpb.SearchResults
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, *querypb.SearchRequest) (*internalpb.SearchResults, error)); ok {
return rf(ctx, req)
}
if rf, ok := ret.Get(0).(func(context.Context, *querypb.SearchRequest) *internalpb.SearchResults); ok {
r0 = rf(ctx, req)
} else {
@ -853,7 +992,6 @@ func (_m *MockQueryNode) Search(ctx context.Context, req *querypb.SearchRequest)
}
}
var r1 error
if rf, ok := ret.Get(1).(func(context.Context, *querypb.SearchRequest) error); ok {
r1 = rf(ctx, req)
} else {
@ -869,8 +1007,8 @@ type MockQueryNode_Search_Call struct {
}
// Search is a helper method to define mock.On call
// - ctx context.Context
// - req *querypb.SearchRequest
// - ctx context.Context
// - req *querypb.SearchRequest
func (_e *MockQueryNode_Expecter) Search(ctx interface{}, req interface{}) *MockQueryNode_Search_Call {
return &MockQueryNode_Search_Call{Call: _e.mock.On("Search", ctx, req)}
}
@ -887,11 +1025,20 @@ func (_c *MockQueryNode_Search_Call) Return(_a0 *internalpb.SearchResults, _a1 e
return _c
}
func (_c *MockQueryNode_Search_Call) RunAndReturn(run func(context.Context, *querypb.SearchRequest) (*internalpb.SearchResults, error)) *MockQueryNode_Search_Call {
_c.Call.Return(run)
return _c
}
// SearchSegments provides a mock function with given fields: ctx, req
func (_m *MockQueryNode) SearchSegments(ctx context.Context, req *querypb.SearchRequest) (*internalpb.SearchResults, error) {
ret := _m.Called(ctx, req)
var r0 *internalpb.SearchResults
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, *querypb.SearchRequest) (*internalpb.SearchResults, error)); ok {
return rf(ctx, req)
}
if rf, ok := ret.Get(0).(func(context.Context, *querypb.SearchRequest) *internalpb.SearchResults); ok {
r0 = rf(ctx, req)
} else {
@ -900,7 +1047,6 @@ func (_m *MockQueryNode) SearchSegments(ctx context.Context, req *querypb.Search
}
}
var r1 error
if rf, ok := ret.Get(1).(func(context.Context, *querypb.SearchRequest) error); ok {
r1 = rf(ctx, req)
} else {
@ -916,8 +1062,8 @@ type MockQueryNode_SearchSegments_Call struct {
}
// SearchSegments is a helper method to define mock.On call
// - ctx context.Context
// - req *querypb.SearchRequest
// - ctx context.Context
// - req *querypb.SearchRequest
func (_e *MockQueryNode_Expecter) SearchSegments(ctx interface{}, req interface{}) *MockQueryNode_SearchSegments_Call {
return &MockQueryNode_SearchSegments_Call{Call: _e.mock.On("SearchSegments", ctx, req)}
}
@ -934,6 +1080,11 @@ func (_c *MockQueryNode_SearchSegments_Call) Return(_a0 *internalpb.SearchResult
return _c
}
func (_c *MockQueryNode_SearchSegments_Call) RunAndReturn(run func(context.Context, *querypb.SearchRequest) (*internalpb.SearchResults, error)) *MockQueryNode_SearchSegments_Call {
_c.Call.Return(run)
return _c
}
// SetAddress provides a mock function with given fields: address
func (_m *MockQueryNode) SetAddress(address string) {
_m.Called(address)
@ -945,7 +1096,7 @@ type MockQueryNode_SetAddress_Call struct {
}
// SetAddress is a helper method to define mock.On call
// - address string
// - address string
func (_e *MockQueryNode_Expecter) SetAddress(address interface{}) *MockQueryNode_SetAddress_Call {
return &MockQueryNode_SetAddress_Call{Call: _e.mock.On("SetAddress", address)}
}
@ -962,6 +1113,11 @@ func (_c *MockQueryNode_SetAddress_Call) Return() *MockQueryNode_SetAddress_Call
return _c
}
func (_c *MockQueryNode_SetAddress_Call) RunAndReturn(run func(string)) *MockQueryNode_SetAddress_Call {
_c.Call.Return(run)
return _c
}
// SetEtcdClient provides a mock function with given fields: etcdClient
func (_m *MockQueryNode) SetEtcdClient(etcdClient *clientv3.Client) {
_m.Called(etcdClient)
@ -973,7 +1129,7 @@ type MockQueryNode_SetEtcdClient_Call struct {
}
// SetEtcdClient is a helper method to define mock.On call
// - etcdClient *clientv3.Client
// - etcdClient *clientv3.Client
func (_e *MockQueryNode_Expecter) SetEtcdClient(etcdClient interface{}) *MockQueryNode_SetEtcdClient_Call {
return &MockQueryNode_SetEtcdClient_Call{Call: _e.mock.On("SetEtcdClient", etcdClient)}
}
@ -990,11 +1146,20 @@ func (_c *MockQueryNode_SetEtcdClient_Call) Return() *MockQueryNode_SetEtcdClien
return _c
}
func (_c *MockQueryNode_SetEtcdClient_Call) RunAndReturn(run func(*clientv3.Client)) *MockQueryNode_SetEtcdClient_Call {
_c.Call.Return(run)
return _c
}
// ShowConfigurations provides a mock function with given fields: ctx, req
func (_m *MockQueryNode) ShowConfigurations(ctx context.Context, req *internalpb.ShowConfigurationsRequest) (*internalpb.ShowConfigurationsResponse, error) {
ret := _m.Called(ctx, req)
var r0 *internalpb.ShowConfigurationsResponse
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, *internalpb.ShowConfigurationsRequest) (*internalpb.ShowConfigurationsResponse, error)); ok {
return rf(ctx, req)
}
if rf, ok := ret.Get(0).(func(context.Context, *internalpb.ShowConfigurationsRequest) *internalpb.ShowConfigurationsResponse); ok {
r0 = rf(ctx, req)
} else {
@ -1003,7 +1168,6 @@ func (_m *MockQueryNode) ShowConfigurations(ctx context.Context, req *internalpb
}
}
var r1 error
if rf, ok := ret.Get(1).(func(context.Context, *internalpb.ShowConfigurationsRequest) error); ok {
r1 = rf(ctx, req)
} else {
@ -1019,8 +1183,8 @@ type MockQueryNode_ShowConfigurations_Call struct {
}
// ShowConfigurations is a helper method to define mock.On call
// - ctx context.Context
// - req *internalpb.ShowConfigurationsRequest
// - ctx context.Context
// - req *internalpb.ShowConfigurationsRequest
func (_e *MockQueryNode_Expecter) ShowConfigurations(ctx interface{}, req interface{}) *MockQueryNode_ShowConfigurations_Call {
return &MockQueryNode_ShowConfigurations_Call{Call: _e.mock.On("ShowConfigurations", ctx, req)}
}
@ -1037,6 +1201,11 @@ func (_c *MockQueryNode_ShowConfigurations_Call) Return(_a0 *internalpb.ShowConf
return _c
}
func (_c *MockQueryNode_ShowConfigurations_Call) RunAndReturn(run func(context.Context, *internalpb.ShowConfigurationsRequest) (*internalpb.ShowConfigurationsResponse, error)) *MockQueryNode_ShowConfigurations_Call {
_c.Call.Return(run)
return _c
}
// Start provides a mock function with given fields:
func (_m *MockQueryNode) Start() error {
ret := _m.Called()
@ -1073,6 +1242,11 @@ func (_c *MockQueryNode_Start_Call) Return(_a0 error) *MockQueryNode_Start_Call
return _c
}
func (_c *MockQueryNode_Start_Call) RunAndReturn(run func() error) *MockQueryNode_Start_Call {
_c.Call.Return(run)
return _c
}
// Stop provides a mock function with given fields:
func (_m *MockQueryNode) Stop() error {
ret := _m.Called()
@ -1109,11 +1283,20 @@ func (_c *MockQueryNode_Stop_Call) Return(_a0 error) *MockQueryNode_Stop_Call {
return _c
}
func (_c *MockQueryNode_Stop_Call) RunAndReturn(run func() error) *MockQueryNode_Stop_Call {
_c.Call.Return(run)
return _c
}
// SyncDistribution provides a mock function with given fields: _a0, _a1
func (_m *MockQueryNode) SyncDistribution(_a0 context.Context, _a1 *querypb.SyncDistributionRequest) (*commonpb.Status, error) {
ret := _m.Called(_a0, _a1)
var r0 *commonpb.Status
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, *querypb.SyncDistributionRequest) (*commonpb.Status, error)); ok {
return rf(_a0, _a1)
}
if rf, ok := ret.Get(0).(func(context.Context, *querypb.SyncDistributionRequest) *commonpb.Status); ok {
r0 = rf(_a0, _a1)
} else {
@ -1122,7 +1305,6 @@ func (_m *MockQueryNode) SyncDistribution(_a0 context.Context, _a1 *querypb.Sync
}
}
var r1 error
if rf, ok := ret.Get(1).(func(context.Context, *querypb.SyncDistributionRequest) error); ok {
r1 = rf(_a0, _a1)
} else {
@ -1138,8 +1320,8 @@ type MockQueryNode_SyncDistribution_Call struct {
}
// SyncDistribution is a helper method to define mock.On call
// - _a0 context.Context
// - _a1 *querypb.SyncDistributionRequest
// - _a0 context.Context
// - _a1 *querypb.SyncDistributionRequest
func (_e *MockQueryNode_Expecter) SyncDistribution(_a0 interface{}, _a1 interface{}) *MockQueryNode_SyncDistribution_Call {
return &MockQueryNode_SyncDistribution_Call{Call: _e.mock.On("SyncDistribution", _a0, _a1)}
}
@ -1156,11 +1338,20 @@ func (_c *MockQueryNode_SyncDistribution_Call) Return(_a0 *commonpb.Status, _a1
return _c
}
func (_c *MockQueryNode_SyncDistribution_Call) RunAndReturn(run func(context.Context, *querypb.SyncDistributionRequest) (*commonpb.Status, error)) *MockQueryNode_SyncDistribution_Call {
_c.Call.Return(run)
return _c
}
// SyncReplicaSegments provides a mock function with given fields: ctx, req
func (_m *MockQueryNode) SyncReplicaSegments(ctx context.Context, req *querypb.SyncReplicaSegmentsRequest) (*commonpb.Status, error) {
ret := _m.Called(ctx, req)
var r0 *commonpb.Status
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, *querypb.SyncReplicaSegmentsRequest) (*commonpb.Status, error)); ok {
return rf(ctx, req)
}
if rf, ok := ret.Get(0).(func(context.Context, *querypb.SyncReplicaSegmentsRequest) *commonpb.Status); ok {
r0 = rf(ctx, req)
} else {
@ -1169,7 +1360,6 @@ func (_m *MockQueryNode) SyncReplicaSegments(ctx context.Context, req *querypb.S
}
}
var r1 error
if rf, ok := ret.Get(1).(func(context.Context, *querypb.SyncReplicaSegmentsRequest) error); ok {
r1 = rf(ctx, req)
} else {
@ -1185,8 +1375,8 @@ type MockQueryNode_SyncReplicaSegments_Call struct {
}
// SyncReplicaSegments is a helper method to define mock.On call
// - ctx context.Context
// - req *querypb.SyncReplicaSegmentsRequest
// - ctx context.Context
// - req *querypb.SyncReplicaSegmentsRequest
func (_e *MockQueryNode_Expecter) SyncReplicaSegments(ctx interface{}, req interface{}) *MockQueryNode_SyncReplicaSegments_Call {
return &MockQueryNode_SyncReplicaSegments_Call{Call: _e.mock.On("SyncReplicaSegments", ctx, req)}
}
@ -1203,11 +1393,20 @@ func (_c *MockQueryNode_SyncReplicaSegments_Call) Return(_a0 *commonpb.Status, _
return _c
}
func (_c *MockQueryNode_SyncReplicaSegments_Call) RunAndReturn(run func(context.Context, *querypb.SyncReplicaSegmentsRequest) (*commonpb.Status, error)) *MockQueryNode_SyncReplicaSegments_Call {
_c.Call.Return(run)
return _c
}
// UnsubDmChannel provides a mock function with given fields: ctx, req
func (_m *MockQueryNode) UnsubDmChannel(ctx context.Context, req *querypb.UnsubDmChannelRequest) (*commonpb.Status, error) {
ret := _m.Called(ctx, req)
var r0 *commonpb.Status
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, *querypb.UnsubDmChannelRequest) (*commonpb.Status, error)); ok {
return rf(ctx, req)
}
if rf, ok := ret.Get(0).(func(context.Context, *querypb.UnsubDmChannelRequest) *commonpb.Status); ok {
r0 = rf(ctx, req)
} else {
@ -1216,7 +1415,6 @@ func (_m *MockQueryNode) UnsubDmChannel(ctx context.Context, req *querypb.UnsubD
}
}
var r1 error
if rf, ok := ret.Get(1).(func(context.Context, *querypb.UnsubDmChannelRequest) error); ok {
r1 = rf(ctx, req)
} else {
@ -1232,8 +1430,8 @@ type MockQueryNode_UnsubDmChannel_Call struct {
}
// UnsubDmChannel is a helper method to define mock.On call
// - ctx context.Context
// - req *querypb.UnsubDmChannelRequest
// - ctx context.Context
// - req *querypb.UnsubDmChannelRequest
func (_e *MockQueryNode_Expecter) UnsubDmChannel(ctx interface{}, req interface{}) *MockQueryNode_UnsubDmChannel_Call {
return &MockQueryNode_UnsubDmChannel_Call{Call: _e.mock.On("UnsubDmChannel", ctx, req)}
}
@ -1250,6 +1448,11 @@ func (_c *MockQueryNode_UnsubDmChannel_Call) Return(_a0 *commonpb.Status, _a1 er
return _c
}
func (_c *MockQueryNode_UnsubDmChannel_Call) RunAndReturn(run func(context.Context, *querypb.UnsubDmChannelRequest) (*commonpb.Status, error)) *MockQueryNode_UnsubDmChannel_Call {
_c.Call.Return(run)
return _c
}
// UpdateStateCode provides a mock function with given fields: stateCode
func (_m *MockQueryNode) UpdateStateCode(stateCode commonpb.StateCode) {
_m.Called(stateCode)
@ -1261,7 +1464,7 @@ type MockQueryNode_UpdateStateCode_Call struct {
}
// UpdateStateCode is a helper method to define mock.On call
// - stateCode commonpb.StateCode
// - stateCode commonpb.StateCode
func (_e *MockQueryNode_Expecter) UpdateStateCode(stateCode interface{}) *MockQueryNode_UpdateStateCode_Call {
return &MockQueryNode_UpdateStateCode_Call{Call: _e.mock.On("UpdateStateCode", stateCode)}
}
@ -1278,11 +1481,20 @@ func (_c *MockQueryNode_UpdateStateCode_Call) Return() *MockQueryNode_UpdateStat
return _c
}
func (_c *MockQueryNode_UpdateStateCode_Call) RunAndReturn(run func(commonpb.StateCode)) *MockQueryNode_UpdateStateCode_Call {
_c.Call.Return(run)
return _c
}
// WatchDmChannels provides a mock function with given fields: ctx, req
func (_m *MockQueryNode) WatchDmChannels(ctx context.Context, req *querypb.WatchDmChannelsRequest) (*commonpb.Status, error) {
ret := _m.Called(ctx, req)
var r0 *commonpb.Status
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, *querypb.WatchDmChannelsRequest) (*commonpb.Status, error)); ok {
return rf(ctx, req)
}
if rf, ok := ret.Get(0).(func(context.Context, *querypb.WatchDmChannelsRequest) *commonpb.Status); ok {
r0 = rf(ctx, req)
} else {
@ -1291,7 +1503,6 @@ func (_m *MockQueryNode) WatchDmChannels(ctx context.Context, req *querypb.Watch
}
}
var r1 error
if rf, ok := ret.Get(1).(func(context.Context, *querypb.WatchDmChannelsRequest) error); ok {
r1 = rf(ctx, req)
} else {
@ -1307,8 +1518,8 @@ type MockQueryNode_WatchDmChannels_Call struct {
}
// WatchDmChannels is a helper method to define mock.On call
// - ctx context.Context
// - req *querypb.WatchDmChannelsRequest
// - ctx context.Context
// - req *querypb.WatchDmChannelsRequest
func (_e *MockQueryNode_Expecter) WatchDmChannels(ctx interface{}, req interface{}) *MockQueryNode_WatchDmChannels_Call {
return &MockQueryNode_WatchDmChannels_Call{Call: _e.mock.On("WatchDmChannels", ctx, req)}
}
@ -1325,13 +1536,17 @@ func (_c *MockQueryNode_WatchDmChannels_Call) Return(_a0 *commonpb.Status, _a1 e
return _c
}
type mockConstructorTestingTNewMockQueryNode interface {
mock.TestingT
Cleanup(func())
func (_c *MockQueryNode_WatchDmChannels_Call) RunAndReturn(run func(context.Context, *querypb.WatchDmChannelsRequest) (*commonpb.Status, error)) *MockQueryNode_WatchDmChannels_Call {
_c.Call.Return(run)
return _c
}
// NewMockQueryNode creates a new instance of MockQueryNode. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
func NewMockQueryNode(t mockConstructorTestingTNewMockQueryNode) *MockQueryNode {
// The first argument is typically a *testing.T value.
func NewMockQueryNode(t interface {
mock.TestingT
Cleanup(func())
}) *MockQueryNode {
mock := &MockQueryNode{}
mock.Mock.Test(t)

499
internal/mocks/mock_rootcoord.go
View File

File diff suppressed because it is too large Load Diff

6
internal/proxy/authentication_interceptor_test.go
View File

@ -7,7 +7,7 @@ import (
"github.com/stretchr/testify/assert"
"google.golang.org/grpc/metadata"
"github.com/milvus-io/milvus/internal/types"
"github.com/milvus-io/milvus/internal/mocks"
"github.com/milvus-io/milvus/pkg/util"
"github.com/milvus-io/milvus/pkg/util/crypto"
"github.com/milvus-io/milvus/pkg/util/paramtable"
@ -32,7 +32,7 @@ func TestValidAuth(t *testing.T) {
assert.False(t, res)
// normal metadata
rootCoord := &MockRootCoordClientInterface{}
queryCoord := &types.MockQueryCoord{}
queryCoord := &mocks.MockQueryCoord{}
mgr := newShardClientMgr()
err := InitMetaCache(ctx, rootCoord, queryCoord, mgr)
assert.NoError(t, err)
@ -65,7 +65,7 @@ func TestAuthenticationInterceptor(t *testing.T) {
assert.Error(t, err)
// mock metacache
rootCoord := &MockRootCoordClientInterface{}
queryCoord := &types.MockQueryCoord{}
queryCoord := &mocks.MockQueryCoord{}
mgr := newShardClientMgr()
err = InitMetaCache(ctx, rootCoord, queryCoord, mgr)
assert.NoError(t, err)

29
internal/proxy/impl_test.go
View File

@ -31,7 +31,6 @@ import (
"github.com/milvus-io/milvus/internal/proto/proxypb"
"github.com/milvus-io/milvus/internal/proto/querypb"
"github.com/milvus-io/milvus/internal/proto/rootcoordpb"
"github.com/milvus-io/milvus/internal/types"
"github.com/milvus-io/milvus/internal/util/dependency"
"github.com/milvus-io/milvus/internal/util/sessionutil"
"github.com/milvus-io/milvus/pkg/log"
@ -76,7 +75,7 @@ func TestProxy_CheckHealth(t *testing.T) {
})
t.Run("proxy health check is ok", func(t *testing.T) {
qc := &types.MockQueryCoord{}
qc := &mocks.MockQueryCoord{}
qc.EXPECT().CheckHealth(mock.Anything, mock.Anything).Return(&milvuspb.CheckHealthResponse{IsHealthy: true}, nil)
node := &Proxy{
rootCoord: NewRootCoordMock(),
@ -105,7 +104,7 @@ func TestProxy_CheckHealth(t *testing.T) {
dataCoordMock := NewDataCoordMock()
dataCoordMock.checkHealthFunc = checkHealthFunc1
qc := &types.MockQueryCoord{}
qc := &mocks.MockQueryCoord{}
qc.EXPECT().CheckHealth(mock.Anything, mock.Anything).Return(nil, errors.New("test"))
node := &Proxy{
session: &sessionutil.Session{ServerID: 1},
@ -124,7 +123,7 @@ func TestProxy_CheckHealth(t *testing.T) {
})
t.Run("check quota state", func(t *testing.T) {
qc := &types.MockQueryCoord{}
qc := &mocks.MockQueryCoord{}
qc.EXPECT().CheckHealth(mock.Anything, mock.Anything).Return(&milvuspb.CheckHealthResponse{IsHealthy: true}, nil)
node := &Proxy{
rootCoord: NewRootCoordMock(),
@ -219,7 +218,7 @@ func TestProxy_ResourceGroup(t *testing.T) {
node.multiRateLimiter = NewMultiRateLimiter()
node.stateCode.Store(commonpb.StateCode_Healthy)
qc := types.NewMockQueryCoord(t)
qc := mocks.NewMockQueryCoord(t)
node.SetQueryCoordClient(qc)
tsoAllocatorIns := newMockTsoAllocator()
@ -311,7 +310,7 @@ func TestProxy_InvalidResourceGroupName(t *testing.T) {
node.multiRateLimiter = NewMultiRateLimiter()
node.stateCode.Store(commonpb.StateCode_Healthy)
qc := types.NewMockQueryCoord(t)
qc := mocks.NewMockQueryCoord(t)
node.SetQueryCoordClient(qc)
qc.EXPECT().DropResourceGroup(mock.Anything, mock.Anything).Return(&commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}, nil)
@ -381,7 +380,7 @@ func TestProxy_FlushAll(t *testing.T) {
err = node.sched.Start()
assert.NoError(t, err)
defer node.sched.Close()
node.dataCoord = mocks.NewDataCoord(t)
node.dataCoord = mocks.NewMockDataCoord(t)
node.rootCoord = mocks.NewRootCoord(t)
// set expectations
@ -399,7 +398,7 @@ func TestProxy_FlushAll(t *testing.T) {
globalMetaCache = cache
successStatus := &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}
node.dataCoord.(*mocks.DataCoord).EXPECT().Flush(mock.Anything, mock.Anything).
node.dataCoord.(*mocks.MockDataCoord).EXPECT().Flush(mock.Anything, mock.Anything).
Return(&datapb.FlushResponse{Status: successStatus}, nil).Maybe()
node.rootCoord.(*mocks.RootCoord).EXPECT().ShowCollections(mock.Anything, mock.Anything).
Return(&milvuspb.ShowCollectionsResponse{Status: successStatus, CollectionNames: []string{"col-0"}}, nil).Maybe()
@ -437,8 +436,8 @@ func TestProxy_FlushAll(t *testing.T) {
})
t.Run("FlushAll failed, DataCoord flush failed", func(t *testing.T) {
node.dataCoord.(*mocks.DataCoord).ExpectedCalls = nil
node.dataCoord.(*mocks.DataCoord).EXPECT().Flush(mock.Anything, mock.Anything).
node.dataCoord.(*mocks.MockDataCoord).ExpectedCalls = nil
node.dataCoord.(*mocks.MockDataCoord).EXPECT().Flush(mock.Anything, mock.Anything).
Return(&datapb.FlushResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
@ -491,12 +490,12 @@ func TestProxy_GetFlushAllState(t *testing.T) {
node.tsoAllocator = &timestampAllocator{
tso: newMockTimestampAllocatorInterface(),
}
node.dataCoord = mocks.NewDataCoord(t)
node.dataCoord = mocks.NewMockDataCoord(t)
node.rootCoord = mocks.NewRootCoord(t)
// set expectations
successStatus := &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}
node.dataCoord.(*mocks.DataCoord).EXPECT().GetFlushAllState(mock.Anything, mock.Anything).
node.dataCoord.(*mocks.MockDataCoord).EXPECT().GetFlushAllState(mock.Anything, mock.Anything).
Return(&milvuspb.GetFlushAllStateResponse{Status: successStatus}, nil).Maybe()
t.Run("GetFlushAllState success", func(t *testing.T) {
@ -514,8 +513,8 @@ func TestProxy_GetFlushAllState(t *testing.T) {
})
t.Run("DataCoord GetFlushAllState failed", func(t *testing.T) {
node.dataCoord.(*mocks.DataCoord).ExpectedCalls = nil
node.dataCoord.(*mocks.DataCoord).EXPECT().GetFlushAllState(mock.Anything, mock.Anything).
node.dataCoord.(*mocks.MockDataCoord).ExpectedCalls = nil
node.dataCoord.(*mocks.MockDataCoord).EXPECT().GetFlushAllState(mock.Anything, mock.Anything).
Return(&milvuspb.GetFlushAllStateResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
@ -538,7 +537,7 @@ func TestProxy_GetReplicas(t *testing.T) {
node.tsoAllocator = &timestampAllocator{
tso: newMockTimestampAllocatorInterface(),
}
mockQC := types.NewMockQueryCoord(t)
mockQC := mocks.NewMockQueryCoord(t)
mockRC := mocks.NewRootCoord(t)
node.queryCoord = mockQC
node.rootCoord = mockRC

11
internal/proxy/lb_policy_test.go
View File

@ -29,6 +29,7 @@ import (
"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-proto/go-api/v2/schemapb"
"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"
@ -41,8 +42,8 @@ import (
type LBPolicySuite struct {
suite.Suite
rc types.RootCoord
qc *types.MockQueryCoord
qn *types.MockQueryNode
qc *mocks.MockQueryCoord
qn *mocks.MockQueryNode
mgr *MockShardClientManager
lbBalancer *MockLBBalancer
@ -50,7 +51,7 @@ type LBPolicySuite struct {
nodes []int64
channels []string
qnList []*types.MockQueryNode
qnList []*mocks.MockQueryNode
collection string
}
@ -63,7 +64,7 @@ func (s *LBPolicySuite) SetupTest() {
s.nodes = []int64{1, 2, 3, 4, 5}
s.channels = []string{"channel1", "channel2"}
successStatus := commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}
qc := types.NewMockQueryCoord(s.T())
qc := mocks.NewMockQueryCoord(s.T())
qc.EXPECT().LoadCollection(mock.Anything, mock.Anything).Return(&successStatus, nil)
qc.EXPECT().GetShardLeaders(mock.Anything, mock.Anything).Return(&querypb.GetShardLeadersResponse{
@ -92,7 +93,7 @@ func (s *LBPolicySuite) SetupTest() {
s.rc = NewRootCoordMock()
s.rc.Start()
s.qn = types.NewMockQueryNode(s.T())
s.qn = mocks.NewMockQueryNode(s.T())
s.qn.EXPECT().GetAddress().Return("localhost").Maybe()
s.qn.EXPECT().GetComponentStates(mock.Anything).Return(nil, nil).Maybe()

8
internal/proxy/look_aside_balancer_test.go
View File

@ -25,8 +25,8 @@ import (
"github.com/cockroachdb/errors"
"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/types"
"github.com/milvus-io/milvus/pkg/util/merr"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/suite"
@ -45,7 +45,7 @@ func (suite *LookAsideBalancerSuite) SetupTest() {
suite.balancer = NewLookAsideBalancer(suite.clientMgr)
suite.balancer.Start(context.Background())
qn := types.NewMockQueryNode(suite.T())
qn := mocks.NewMockQueryNode(suite.T())
suite.clientMgr.EXPECT().GetClient(mock.Anything, int64(1)).Return(qn, nil).Maybe()
qn.EXPECT().GetComponentStates(mock.Anything).Return(nil, errors.New("fake error")).Maybe()
}
@ -289,7 +289,7 @@ func (suite *LookAsideBalancerSuite) TestCancelWorkload() {
}
func (suite *LookAsideBalancerSuite) TestCheckHealthLoop() {
qn2 := types.NewMockQueryNode(suite.T())
qn2 := mocks.NewMockQueryNode(suite.T())
suite.clientMgr.EXPECT().GetClient(mock.Anything, int64(2)).Return(qn2, nil)
qn2.EXPECT().GetComponentStates(mock.Anything).Return(&milvuspb.ComponentStates{
State: &milvuspb.ComponentInfo{
@ -314,7 +314,7 @@ func (suite *LookAsideBalancerSuite) TestCheckHealthLoop() {
func (suite *LookAsideBalancerSuite) TestNodeRecover() {
// mock qn down for a while and then recover
qn3 := types.NewMockQueryNode(suite.T())
qn3 := mocks.NewMockQueryNode(suite.T())
suite.clientMgr.EXPECT().GetClient(mock.Anything, int64(3)).Return(qn3, nil)
qn3.EXPECT().GetComponentStates(mock.Anything).Return(&milvuspb.ComponentStates{
State: &milvuspb.ComponentInfo{

21
internal/proxy/meta_cache_test.go
View File

@ -33,6 +33,7 @@ import (
"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-proto/go-api/v2/schemapb"
"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/proto/rootcoordpb"
@ -213,7 +214,7 @@ func (m *MockRootCoordClientInterface) ListPolicy(ctx context.Context, in *inter
func TestMetaCache_GetCollection(t *testing.T) {
ctx := context.Background()
rootCoord := &MockRootCoordClientInterface{}
queryCoord := &types.MockQueryCoord{}
queryCoord := &mocks.MockQueryCoord{}
mgr := newShardClientMgr()
err := InitMetaCache(ctx, rootCoord, queryCoord, mgr)
assert.NoError(t, err)
@ -264,7 +265,7 @@ func TestMetaCache_GetCollection(t *testing.T) {
func TestMetaCache_GetCollectionName(t *testing.T) {
ctx := context.Background()
rootCoord := &MockRootCoordClientInterface{}
queryCoord := &types.MockQueryCoord{}
queryCoord := &mocks.MockQueryCoord{}
mgr := newShardClientMgr()
err := InitMetaCache(ctx, rootCoord, queryCoord, mgr)
assert.NoError(t, err)
@ -315,7 +316,7 @@ func TestMetaCache_GetCollectionName(t *testing.T) {
func TestMetaCache_GetCollectionFailure(t *testing.T) {
ctx := context.Background()
rootCoord := &MockRootCoordClientInterface{}
queryCoord := &types.MockQueryCoord{}
queryCoord := &mocks.MockQueryCoord{}
mgr := newShardClientMgr()
err := InitMetaCache(ctx, rootCoord, queryCoord, mgr)
assert.NoError(t, err)
@ -348,7 +349,7 @@ func TestMetaCache_GetCollectionFailure(t *testing.T) {
func TestMetaCache_GetNonExistCollection(t *testing.T) {
ctx := context.Background()
rootCoord := &MockRootCoordClientInterface{}
queryCoord := &types.MockQueryCoord{}
queryCoord := &mocks.MockQueryCoord{}
mgr := newShardClientMgr()
err := InitMetaCache(ctx, rootCoord, queryCoord, mgr)
assert.NoError(t, err)
@ -364,7 +365,7 @@ func TestMetaCache_GetNonExistCollection(t *testing.T) {
func TestMetaCache_GetPartitionID(t *testing.T) {
ctx := context.Background()
rootCoord := &MockRootCoordClientInterface{}
queryCoord := &types.MockQueryCoord{}
queryCoord := &mocks.MockQueryCoord{}
mgr := newShardClientMgr()
err := InitMetaCache(ctx, rootCoord, queryCoord, mgr)
assert.NoError(t, err)
@ -386,7 +387,7 @@ func TestMetaCache_GetPartitionID(t *testing.T) {
func TestMetaCache_ConcurrentTest1(t *testing.T) {
ctx := context.Background()
rootCoord := &MockRootCoordClientInterface{}
queryCoord := &types.MockQueryCoord{}
queryCoord := &mocks.MockQueryCoord{}
mgr := newShardClientMgr()
err := InitMetaCache(ctx, rootCoord, queryCoord, mgr)
assert.NoError(t, err)
@ -440,7 +441,7 @@ func TestMetaCache_ConcurrentTest1(t *testing.T) {
func TestMetaCache_GetPartitionError(t *testing.T) {
ctx := context.Background()
rootCoord := &MockRootCoordClientInterface{}
queryCoord := &types.MockQueryCoord{}
queryCoord := &mocks.MockQueryCoord{}
mgr := newShardClientMgr()
err := InitMetaCache(ctx, rootCoord, queryCoord, mgr)
assert.NoError(t, err)
@ -609,7 +610,7 @@ func TestMetaCache_ClearShards(t *testing.T) {
func TestMetaCache_PolicyInfo(t *testing.T) {
client := &MockRootCoordClientInterface{}
qc := &types.MockQueryCoord{}
qc := &mocks.MockQueryCoord{}
mgr := newShardClientMgr()
t.Run("InitMetaCache", func(t *testing.T) {
@ -692,7 +693,7 @@ func TestMetaCache_PolicyInfo(t *testing.T) {
func TestMetaCache_RemoveCollection(t *testing.T) {
ctx := context.Background()
rootCoord := &MockRootCoordClientInterface{}
queryCoord := &types.MockQueryCoord{}
queryCoord := &mocks.MockQueryCoord{}
shardMgr := newShardClientMgr()
err := InitMetaCache(ctx, rootCoord, queryCoord, shardMgr)
assert.NoError(t, err)
@ -736,7 +737,7 @@ func TestMetaCache_ExpireShardLeaderCache(t *testing.T) {
ctx := context.Background()
rootCoord := &MockRootCoordClientInterface{}
queryCoord := &types.MockQueryCoord{}
queryCoord := &mocks.MockQueryCoord{}
shardMgr := newShardClientMgr()
err := InitMetaCache(ctx, rootCoord, queryCoord, shardMgr)
assert.NoError(t, err)

6
internal/proxy/privilege_interceptor_test.go
View File

@ -7,8 +7,8 @@ import (
"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/types"
"github.com/milvus-io/milvus/pkg/util/funcutil"
"github.com/milvus-io/milvus/pkg/util/paramtable"
"github.com/stretchr/testify/assert"
@ -45,7 +45,7 @@ func TestPrivilegeInterceptor(t *testing.T) {
ctx = GetContext(context.Background(), "alice:123456")
client := &MockRootCoordClientInterface{}
queryCoord := &types.MockQueryCoord{}
queryCoord := &mocks.MockQueryCoord{}
mgr := newShardClientMgr()
client.listPolicy = func(ctx context.Context, in *internalpb.ListPolicyRequest) (*internalpb.ListPolicyResponse, error) {
@ -176,7 +176,7 @@ func TestResourceGroupPrivilege(t *testing.T) {
ctx = GetContext(context.Background(), "fooo:123456")
client := &MockRootCoordClientInterface{}
queryCoord := &types.MockQueryCoord{}
queryCoord := &mocks.MockQueryCoord{}
mgr := newShardClientMgr()
client.listPolicy = func(ctx context.Context, in *internalpb.ListPolicyRequest) (*internalpb.ListPolicyResponse, error) {

6
internal/proxy/task_index_test.go
View File

@ -29,9 +29,9 @@ import (
"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-proto/go-api/v2/schemapb"
"github.com/milvus-io/milvus/internal/mocks"
"github.com/milvus-io/milvus/internal/proto/indexpb"
"github.com/milvus-io/milvus/internal/proto/querypb"
"github.com/milvus-io/milvus/internal/types"
"github.com/milvus-io/milvus/pkg/common"
"github.com/milvus-io/milvus/pkg/config"
"github.com/milvus-io/milvus/pkg/util/funcutil"
@ -218,8 +218,8 @@ func TestDropIndexTask_PreExecute(t *testing.T) {
})
}
func getMockQueryCoord() *types.MockQueryCoord {
qc := &types.MockQueryCoord{}
func getMockQueryCoord() *mocks.MockQueryCoord {
qc := &mocks.MockQueryCoord{}
successStatus := &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}
qc.EXPECT().LoadCollection(mock.Anything, mock.Anything).Return(successStatus, nil)
qc.EXPECT().GetShardLeaders(mock.Anything, mock.Anything).Return(&querypb.GetShardLeadersResponse{

4
internal/proxy/task_query_test.go
View File

@ -32,9 +32,9 @@ import (
"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-proto/go-api/v2/schemapb"
"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/common"
"github.com/milvus-io/milvus/pkg/log"
"github.com/milvus-io/milvus/pkg/util/funcutil"
@ -48,7 +48,7 @@ func TestQueryTask_all(t *testing.T) {
ctx = context.TODO()
rc = NewRootCoordMock()
qc = types.NewMockQueryCoord(t)
qc = mocks.NewMockQueryCoord(t)
qn = getQueryNode()
shardsNum = common.DefaultShardsNum

24
internal/proxy/task_search_test.go
View File

@ -152,7 +152,7 @@ func TestSearchTask_PreExecute(t *testing.T) {
var (
rc = NewRootCoordMock()
qc = types.NewMockQueryCoord(t)
qc = mocks.NewMockQueryCoord(t)
ctx = context.TODO()
)
@ -258,15 +258,15 @@ func TestSearchTask_PreExecute(t *testing.T) {
})
}
func getQueryCoord() *types.MockQueryCoord {
qc := &types.MockQueryCoord{}
func getQueryCoord() *mocks.MockQueryCoord {
qc := &mocks.MockQueryCoord{}
qc.EXPECT().Start().Return(nil)
qc.EXPECT().Stop().Return(nil)
return qc
}
func getQueryNode() *types.MockQueryNode {
qn := &types.MockQueryNode{}
func getQueryNode() *mocks.MockQueryNode {
qn := &mocks.MockQueryNode{}
return qn
}
@ -1868,7 +1868,7 @@ func TestSearchTask_Requery(t *testing.T) {
t.Run("Test normal", func(t *testing.T) {
schema := constructCollectionSchema(pkField, vecField, dim, collection)
node := mocks.NewProxy(t)
node := mocks.NewMockProxy(t)
node.EXPECT().Query(mock.Anything, mock.Anything).
Return(&milvuspb.QueryResults{
FieldsData: []*schemapb.FieldData{{
@ -1926,7 +1926,7 @@ func TestSearchTask_Requery(t *testing.T) {
t.Run("Test no primary key", func(t *testing.T) {
schema := &schemapb.CollectionSchema{}
node := mocks.NewProxy(t)
node := mocks.NewMockProxy(t)
qt := &searchTask{
ctx: ctx,
@ -1949,7 +1949,7 @@ func TestSearchTask_Requery(t *testing.T) {
t.Run("Test requery failed 1", func(t *testing.T) {
schema := constructCollectionSchema(pkField, vecField, dim, collection)
node := mocks.NewProxy(t)
node := mocks.NewMockProxy(t)
node.EXPECT().Query(mock.Anything, mock.Anything).
Return(nil, fmt.Errorf("mock err 1"))
@ -1974,7 +1974,7 @@ func TestSearchTask_Requery(t *testing.T) {
t.Run("Test requery failed 2", func(t *testing.T) {
schema := constructCollectionSchema(pkField, vecField, dim, collection)
node := mocks.NewProxy(t)
node := mocks.NewMockProxy(t)
node.EXPECT().Query(mock.Anything, mock.Anything).
Return(&milvuspb.QueryResults{
Status: &commonpb.Status{
@ -2004,7 +2004,7 @@ func TestSearchTask_Requery(t *testing.T) {
t.Run("Test get pk field data failed", func(t *testing.T) {
schema := constructCollectionSchema(pkField, vecField, dim, collection)
node := mocks.NewProxy(t)
node := mocks.NewMockProxy(t)
node.EXPECT().Query(mock.Anything, mock.Anything).
Return(&milvuspb.QueryResults{
FieldsData: []*schemapb.FieldData{},
@ -2031,7 +2031,7 @@ func TestSearchTask_Requery(t *testing.T) {
t.Run("Test incomplete query result", func(t *testing.T) {
schema := constructCollectionSchema(pkField, vecField, dim, collection)
node := mocks.NewProxy(t)
node := mocks.NewMockProxy(t)
node.EXPECT().Query(mock.Anything, mock.Anything).
Return(&milvuspb.QueryResults{
FieldsData: []*schemapb.FieldData{{
@ -2085,7 +2085,7 @@ func TestSearchTask_Requery(t *testing.T) {
t.Run("Test postExecute with requery failed", func(t *testing.T) {
schema := constructCollectionSchema(pkField, vecField, dim, collection)
node := mocks.NewProxy(t)
node := mocks.NewMockProxy(t)
node.EXPECT().Query(mock.Anything, mock.Anything).
Return(nil, fmt.Errorf("mock err 1"))

7
internal/proxy/task_statistic_test.go
View File

@ -27,6 +27,7 @@ import (
"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-proto/go-api/v2/schemapb"
"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"
@ -39,7 +40,7 @@ type StatisticTaskSuite struct {
suite.Suite
rc types.RootCoord
qc types.QueryCoord
qn *types.MockQueryNode
qn *mocks.MockQueryNode
lb LBPolicy
@ -52,7 +53,7 @@ func (s *StatisticTaskSuite) SetupSuite() {
func (s *StatisticTaskSuite) SetupTest() {
successStatus := commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}
qc := types.NewMockQueryCoord(s.T())
qc := mocks.NewMockQueryCoord(s.T())
qc.EXPECT().LoadCollection(mock.Anything, mock.Anything).Return(&successStatus, nil)
qc.EXPECT().GetShardLeaders(mock.Anything, mock.Anything).Return(&querypb.GetShardLeadersResponse{
@ -75,7 +76,7 @@ func (s *StatisticTaskSuite) SetupTest() {
s.qc = qc
s.rc = NewRootCoordMock()
s.rc.Start()
s.qn = types.NewMockQueryNode(s.T())
s.qn = mocks.NewMockQueryNode(s.T())
s.qn.EXPECT().GetComponentStates(mock.Anything).Return(nil, nil).Maybe()
mgr := NewMockShardClientManager(s.T())

16
internal/proxy/util_test.go
View File

@ -33,10 +33,10 @@ import (
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
"github.com/milvus-io/milvus-proto/go-api/v2/msgpb"
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
"github.com/milvus-io/milvus/internal/mocks"
"github.com/milvus-io/milvus/internal/proto/internalpb"
"github.com/milvus-io/milvus/internal/proto/querypb"
"github.com/milvus-io/milvus/internal/proto/rootcoordpb"
"github.com/milvus-io/milvus/internal/types"
"github.com/milvus-io/milvus/pkg/common"
"github.com/milvus-io/milvus/pkg/util"
"github.com/milvus-io/milvus/pkg/util/crypto"
@ -900,7 +900,7 @@ func Test_isCollectionIsLoaded(t *testing.T) {
ctx := context.Background()
t.Run("normal", func(t *testing.T) {
collID := int64(1)
qc := &types.MockQueryCoord{}
qc := &mocks.MockQueryCoord{}
successStatus := &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}
qc.EXPECT().LoadCollection(mock.Anything, mock.Anything).Return(successStatus, nil)
qc.EXPECT().GetShardLeaders(mock.Anything, mock.Anything).Return(&querypb.GetShardLeadersResponse{
@ -924,7 +924,7 @@ func Test_isCollectionIsLoaded(t *testing.T) {
t.Run("error", func(t *testing.T) {
collID := int64(1)
qc := &types.MockQueryCoord{}
qc := &mocks.MockQueryCoord{}
successStatus := &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}
qc.EXPECT().LoadCollection(mock.Anything, mock.Anything).Return(successStatus, nil)
qc.EXPECT().GetShardLeaders(mock.Anything, mock.Anything).Return(&querypb.GetShardLeadersResponse{
@ -948,7 +948,7 @@ func Test_isCollectionIsLoaded(t *testing.T) {
t.Run("fail", func(t *testing.T) {
collID := int64(1)
qc := &types.MockQueryCoord{}
qc := &mocks.MockQueryCoord{}
successStatus := &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}
qc.EXPECT().LoadCollection(mock.Anything, mock.Anything).Return(successStatus, nil)
qc.EXPECT().GetShardLeaders(mock.Anything, mock.Anything).Return(&querypb.GetShardLeadersResponse{
@ -979,7 +979,7 @@ func Test_isPartitionIsLoaded(t *testing.T) {
t.Run("normal", func(t *testing.T) {
collID := int64(1)
partID := int64(2)
qc := &types.MockQueryCoord{}
qc := &mocks.MockQueryCoord{}
successStatus := &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}
qc.EXPECT().LoadCollection(mock.Anything, mock.Anything).Return(successStatus, nil)
qc.EXPECT().GetShardLeaders(mock.Anything, mock.Anything).Return(&querypb.GetShardLeadersResponse{
@ -1007,7 +1007,7 @@ func Test_isPartitionIsLoaded(t *testing.T) {
t.Run("error", func(t *testing.T) {
collID := int64(1)
partID := int64(2)
qc := &types.MockQueryCoord{}
qc := &mocks.MockQueryCoord{}
successStatus := &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}
qc.EXPECT().LoadCollection(mock.Anything, mock.Anything).Return(successStatus, nil)
qc.EXPECT().GetShardLeaders(mock.Anything, mock.Anything).Return(&querypb.GetShardLeadersResponse{
@ -1035,7 +1035,7 @@ func Test_isPartitionIsLoaded(t *testing.T) {
t.Run("fail", func(t *testing.T) {
collID := int64(1)
partID := int64(2)
qc := &types.MockQueryCoord{}
qc := &mocks.MockQueryCoord{}
successStatus := &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}
qc.EXPECT().LoadCollection(mock.Anything, mock.Anything).Return(successStatus, nil)
qc.EXPECT().GetShardLeaders(mock.Anything, mock.Anything).Return(&querypb.GetShardLeadersResponse{
@ -1833,7 +1833,7 @@ func Test_MaxQueryResultWindow(t *testing.T) {
}
func Test_GetPartitionProgressFailed(t *testing.T) {
qc := types.NewMockQueryCoord(t)
qc := mocks.NewMockQueryCoord(t)
qc.EXPECT().ShowPartitions(mock.Anything, mock.Anything).Return(&querypb.ShowPartitionsResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,

6
internal/querycoordv2/meta/coordinator_broker_test.go
View File

@ -78,7 +78,7 @@ func TestCoordinatorBroker_GetCollectionSchema(t *testing.T) {
func TestCoordinatorBroker_GetRecoveryInfo(t *testing.T) {
t.Run("normal case", func(t *testing.T) {
dc := mocks.NewDataCoord(t)
dc := mocks.NewMockDataCoord(t)
dc.EXPECT().GetRecoveryInfoV2(mock.Anything, mock.Anything).Return(&datapb.GetRecoveryInfoResponseV2{}, nil)
ctx := context.Background()
@ -89,7 +89,7 @@ func TestCoordinatorBroker_GetRecoveryInfo(t *testing.T) {
})
t.Run("get error", func(t *testing.T) {
dc := mocks.NewDataCoord(t)
dc := mocks.NewMockDataCoord(t)
fakeErr := errors.New("fake error")
dc.EXPECT().GetRecoveryInfoV2(mock.Anything, mock.Anything).Return(nil, fakeErr)
@ -101,7 +101,7 @@ func TestCoordinatorBroker_GetRecoveryInfo(t *testing.T) {
})
t.Run("return non-success code", func(t *testing.T) {
dc := mocks.NewDataCoord(t)
dc := mocks.NewMockDataCoord(t)
dc.EXPECT().GetRecoveryInfoV2(mock.Anything, mock.Anything).Return(&datapb.GetRecoveryInfoResponseV2{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,

6
internal/querycoordv2/server_test.go
View File

@ -296,7 +296,7 @@ func (suite *ServerSuite) TestEnableActiveStandby() {
suite.server, err = suite.newQueryCoord()
suite.NoError(err)
mockRootCoord := coordMocks.NewRootCoord(suite.T())
mockDataCoord := coordMocks.NewDataCoord(suite.T())
mockDataCoord := coordMocks.NewMockDataCoord(suite.T())
mockRootCoord.EXPECT().DescribeCollection(mock.Anything, mock.Anything).Return(&milvuspb.DescribeCollectionResponse{
Status: merr.Status(nil),
@ -415,7 +415,7 @@ func (suite *ServerSuite) expectLoadAndReleasePartitions(querynode *mocks.MockQu
querynode.EXPECT().ReleasePartitions(mock.Anything, mock.Anything).Return(utils.WrapStatus(commonpb.ErrorCode_Success, ""), nil).Maybe()
}
func (suite *ServerSuite) expectGetRecoverInfoByMockDataCoord(collection int64, dataCoord *coordMocks.DataCoord) {
func (suite *ServerSuite) expectGetRecoverInfoByMockDataCoord(collection int64, dataCoord *coordMocks.MockDataCoord) {
var (
vChannels []*datapb.VchannelInfo
segmentInfos []*datapb.SegmentInfo
@ -532,7 +532,7 @@ func (suite *ServerSuite) hackServer() {
func (suite *ServerSuite) hackBroker(server *Server) {
mockRootCoord := coordMocks.NewRootCoord(suite.T())
mockDataCoord := coordMocks.NewDataCoord(suite.T())
mockDataCoord := coordMocks.NewMockDataCoord(suite.T())
for _, collection := range suite.collections {
mockRootCoord.EXPECT().DescribeCollection(mock.Anything, mock.Anything).Return(&milvuspb.DescribeCollectionResponse{

8
internal/querynodev2/cluster/worker_test.go
View File

@ -27,9 +27,9 @@ import (
"google.golang.org/grpc/status"
commonpb "github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus/internal/mocks"
internalpb "github.com/milvus-io/milvus/internal/proto/internalpb"
querypb "github.com/milvus-io/milvus/internal/proto/querypb"
"github.com/milvus-io/milvus/internal/types"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/suite"
@ -38,12 +38,12 @@ import (
type RemoteWorkerSuite struct {
suite.Suite
mockClient *types.MockQueryNode
mockClient *mocks.MockQueryNode
worker *remoteWorker
}
func (s *RemoteWorkerSuite) SetupTest() {
s.mockClient = &types.MockQueryNode{}
s.mockClient = &mocks.MockQueryNode{}
s.worker = &remoteWorker{client: s.mockClient}
}
@ -404,7 +404,7 @@ func TestRemoteWorker(t *testing.T) {
}
func TestNewRemoteWorker(t *testing.T) {
client := &types.MockQueryNode{}
client := &mocks.MockQueryNode{}
w := NewRemoteWorker(client)

6
internal/rootcoord/broker_test.go
View File

@ -342,7 +342,7 @@ func TestServerBroker_BroadcastAlteredCollection(t *testing.T) {
func TestServerBroker_GcConfirm(t *testing.T) {
t.Run("invalid datacoord", func(t *testing.T) {
dc := mocks.NewDataCoord(t)
dc := mocks.NewMockDataCoord(t)
dc.On("GcConfirm",
mock.Anything, // context.Context
mock.Anything, // *datapb.GcConfirmRequest
@ -353,7 +353,7 @@ func TestServerBroker_GcConfirm(t *testing.T) {
})
t.Run("non success", func(t *testing.T) {
dc := mocks.NewDataCoord(t)
dc := mocks.NewMockDataCoord(t)
dc.On("GcConfirm",
mock.Anything, // context.Context
mock.Anything, // *datapb.GcConfirmRequest
@ -366,7 +366,7 @@ func TestServerBroker_GcConfirm(t *testing.T) {
})
t.Run("normal case", func(t *testing.T) {
dc := mocks.NewDataCoord(t)
dc := mocks.NewMockDataCoord(t)
dc.On("GcConfirm",
mock.Anything, // context.Context
mock.Anything, // *datapb.GcConfirmRequest

8
internal/rootcoord/garbage_collector_test.go
View File

@ -26,10 +26,10 @@ import (
"github.com/stretchr/testify/mock"
"github.com/milvus-io/milvus/internal/metastore/model"
"github.com/milvus-io/milvus/internal/mocks"
"github.com/milvus-io/milvus/internal/proto/querypb"
mockrootcoord "github.com/milvus-io/milvus/internal/rootcoord/mocks"
mocktso "github.com/milvus-io/milvus/internal/tso/mocks"
"github.com/milvus-io/milvus/internal/types"
"github.com/milvus-io/milvus/pkg/common"
"github.com/milvus-io/milvus/pkg/util/merr"
)
@ -460,7 +460,7 @@ func TestGarbageCollector_RemoveCreatingPartition(t *testing.T) {
signal <- struct{}{}
})
qc := types.NewMockQueryCoord(t)
qc := mocks.NewMockQueryCoord(t)
qc.EXPECT().ReleasePartitions(mock.Anything, mock.Anything).Return(merr.Status(nil), nil)
core := newTestCore(withTtSynchronizer(ticker),
@ -485,7 +485,7 @@ func TestGarbageCollector_RemoveCreatingPartition(t *testing.T) {
signal := make(chan struct{}, 1)
meta := mockrootcoord.NewIMetaTable(t)
qc := types.NewMockQueryCoord(t)
qc := mocks.NewMockQueryCoord(t)
qc.EXPECT().ReleasePartitions(mock.Anything, mock.Anything).
Return(merr.Status(nil), fmt.Errorf("mock err")).
Run(func(ctx context.Context, req *querypb.ReleasePartitionsRequest) {
@ -519,7 +519,7 @@ func TestGarbageCollector_RemoveCreatingPartition(t *testing.T) {
signal <- struct{}{}
})
qc := types.NewMockQueryCoord(t)
qc := mocks.NewMockQueryCoord(t)
qc.EXPECT().ReleasePartitions(mock.Anything, mock.Anything).Return(merr.Status(nil), nil)
core := newTestCore(withTtSynchronizer(ticker),

9
internal/rootcoord/mock_test.go
View File

@ -29,6 +29,7 @@ import (
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
"github.com/milvus-io/milvus/internal/allocator"
"github.com/milvus-io/milvus/internal/metastore/model"
"github.com/milvus-io/milvus/internal/mocks"
"github.com/milvus-io/milvus/internal/proto/datapb"
pb "github.com/milvus-io/milvus/internal/proto/etcdpb"
"github.com/milvus-io/milvus/internal/proto/indexpb"
@ -445,7 +446,7 @@ func withQueryCoord(qc types.QueryCoord) Opt {
}
func withUnhealthyQueryCoord() Opt {
qc := &types.MockQueryCoord{}
qc := &mocks.MockQueryCoord{}
qc.EXPECT().GetComponentStates(mock.Anything).Return(
&milvuspb.ComponentStates{
State: &milvuspb.ComponentInfo{StateCode: commonpb.StateCode_Abnormal},
@ -456,7 +457,7 @@ func withUnhealthyQueryCoord() Opt {
}
func withInvalidQueryCoord() Opt {
qc := &types.MockQueryCoord{}
qc := &mocks.MockQueryCoord{}
qc.EXPECT().GetComponentStates(mock.Anything).Return(
&milvuspb.ComponentStates{
State: &milvuspb.ComponentInfo{StateCode: commonpb.StateCode_Healthy},
@ -475,7 +476,7 @@ func withInvalidQueryCoord() Opt {
}
func withFailedQueryCoord() Opt {
qc := &types.MockQueryCoord{}
qc := &mocks.MockQueryCoord{}
qc.EXPECT().GetComponentStates(mock.Anything).Return(
&milvuspb.ComponentStates{
State: &milvuspb.ComponentInfo{StateCode: commonpb.StateCode_Healthy},
@ -496,7 +497,7 @@ func withFailedQueryCoord() Opt {
}
func withValidQueryCoord() Opt {
qc := &types.MockQueryCoord{}
qc := &mocks.MockQueryCoord{}
qc.EXPECT().GetComponentStates(mock.Anything).Return(
&milvuspb.ComponentStates{
State: &milvuspb.ComponentInfo{StateCode: commonpb.StateCode_Healthy},

31
internal/rootcoord/quota_center_test.go
View File

@ -29,6 +29,7 @@ import (
"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/metastore/model"
"github.com/milvus-io/milvus/internal/mocks"
"github.com/milvus-io/milvus/internal/proto/internalpb"
mockrootcoord "github.com/milvus-io/milvus/internal/rootcoord/mocks"
"github.com/milvus-io/milvus/internal/types"
@ -72,7 +73,7 @@ func TestQuotaCenter(t *testing.T) {
pcm := newProxyClientManager(core.proxyCreator)
t.Run("test QuotaCenter", func(t *testing.T) {
qc := types.NewMockQueryCoord(t)
qc := mocks.NewMockQueryCoord(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)
@ -82,7 +83,7 @@ func TestQuotaCenter(t *testing.T) {
})
t.Run("test syncMetrics", func(t *testing.T) {
qc := types.NewMockQueryCoord(t)
qc := mocks.NewMockQueryCoord(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: succStatus()}, nil)
@ -112,7 +113,7 @@ func TestQuotaCenter(t *testing.T) {
})
t.Run("test forceDeny", func(t *testing.T) {
qc := types.NewMockQueryCoord(t)
qc := mocks.NewMockQueryCoord(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)
@ -142,7 +143,7 @@ func TestQuotaCenter(t *testing.T) {
})
t.Run("test calculateRates", func(t *testing.T) {
qc := types.NewMockQueryCoord(t)
qc := mocks.NewMockQueryCoord(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)
@ -158,7 +159,7 @@ func TestQuotaCenter(t *testing.T) {
})
t.Run("test getTimeTickDelayFactor factors", func(t *testing.T) {
qc := types.NewMockQueryCoord(t)
qc := mocks.NewMockQueryCoord(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)
@ -206,7 +207,7 @@ func TestQuotaCenter(t *testing.T) {
})
t.Run("test TimeTickDelayFactor factors", func(t *testing.T) {
qc := types.NewMockQueryCoord(t)
qc := mocks.NewMockQueryCoord(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)
@ -279,7 +280,7 @@ func TestQuotaCenter(t *testing.T) {
})
t.Run("test calculateReadRates", func(t *testing.T) {
qc := types.NewMockQueryCoord(t)
qc := mocks.NewMockQueryCoord(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)
@ -338,7 +339,7 @@ func TestQuotaCenter(t *testing.T) {
})
t.Run("test calculateWriteRates", func(t *testing.T) {
qc := types.NewMockQueryCoord(t)
qc := mocks.NewMockQueryCoord(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)
@ -378,7 +379,7 @@ func TestQuotaCenter(t *testing.T) {
})
t.Run("test MemoryFactor factors", func(t *testing.T) {
qc := types.NewMockQueryCoord(t)
qc := mocks.NewMockQueryCoord(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)
@ -433,7 +434,7 @@ func TestQuotaCenter(t *testing.T) {
})
t.Run("test GrowingSegmentsSize factors", func(t *testing.T) {
qc := types.NewMockQueryCoord(t)
qc := mocks.NewMockQueryCoord(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)
@ -487,7 +488,7 @@ func TestQuotaCenter(t *testing.T) {
})
t.Run("test checkDiskQuota", func(t *testing.T) {
qc := types.NewMockQueryCoord(t)
qc := mocks.NewMockQueryCoord(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)
@ -530,8 +531,8 @@ func TestQuotaCenter(t *testing.T) {
})
t.Run("test setRates", func(t *testing.T) {
qc := types.NewMockQueryCoord(t)
p1 := types.NewMockProxy(t)
qc := mocks.NewMockQueryCoord(t)
p1 := mocks.NewMockProxy(t)
p1.EXPECT().SetRates(mock.Anything, mock.Anything).Return(nil, nil)
pcm := &proxyClientManager{proxyClient: map[int64]types.Proxy{
TestProxyID: p1,
@ -551,7 +552,7 @@ func TestQuotaCenter(t *testing.T) {
})
t.Run("test recordMetrics", func(t *testing.T) {
qc := types.NewMockQueryCoord(t)
qc := mocks.NewMockQueryCoord(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)
@ -563,7 +564,7 @@ func TestQuotaCenter(t *testing.T) {
})
t.Run("test guaranteeMinRate", func(t *testing.T) {
qc := types.NewMockQueryCoord(t)
qc := mocks.NewMockQueryCoord(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)

2793
internal/types/mock_datacoord.go
View File

File diff suppressed because it is too large Load Diff

4225
internal/types/mock_proxy.go
View File

File diff suppressed because it is too large Load Diff

10
internal/types/types.go
View File

@ -102,6 +102,8 @@ type DataNode interface {
}
// DataNodeComponent is used by grpc server of DataNode
//
//go:generate mockery --name=DataNodeComponent --structname=MockDataNode --output=../mocks --filename=mock_datanode.go --with-expecter
type DataNodeComponent interface {
DataNode
@ -381,6 +383,8 @@ type DataCoord interface {
}
// DataCoordComponent defines the interface of DataCoord component.
//
//go:generate mockery --name=DataCoordComponent --structname=MockDataCoord --output=../mocks --filename=mock_datacoord.go --with-expecter
type DataCoordComponent interface {
DataCoord
@ -885,6 +889,8 @@ type Proxy interface {
}
// ProxyComponent defines the interface of proxy component.
//
//go:generate mockery --name=ProxyComponent --structname=MockProxy --output=../mocks --filename=mock_proxy.go --with-expecter
type ProxyComponent interface {
Proxy
@ -1452,6 +1458,8 @@ type QueryNode interface {
}
// QueryNodeComponent is used by grpc server of QueryNode
//
//go:generate mockery --name=QueryNodeComponent --structname=MockQueryNode --output=../mocks --filename=mock_querynode.go --with-expecter
type QueryNodeComponent interface {
QueryNode
@ -1499,6 +1507,8 @@ type QueryCoord interface {
}
// QueryCoordComponent is used by grpc server of QueryCoord
//
//go:generate mockery --name=QueryCoordComponent --structname=MockQueryCoord --output=../mocks --filename=mock_querycoord.go --with-expecter
type QueryCoordComponent interface {
QueryCoord