milvus/internal/querycoordv2/meta/mock_target_manager.go
congqixia 4aa8a12ce8
fix: [2.4] Check partition in current target when observing partition load status (#34282) (#34305)
Cherry-pick from master
pr: #34282
See also #34234

`LoadPartitions` does not guarantee the current target has loading
partitions if there are some partitions already loaded before.

This PR check current target contains the partition to load when
advancing loading percentage to 100.

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
2024-07-02 15:48:10 +08:00

977 lines
36 KiB
Go

// Code generated by mockery v2.32.4. DO NOT EDIT.
package meta
import (
metastore "github.com/milvus-io/milvus/internal/metastore"
datapb "github.com/milvus-io/milvus/internal/proto/datapb"
mock "github.com/stretchr/testify/mock"
typeutil "github.com/milvus-io/milvus/pkg/util/typeutil"
)
// MockTargetManager is an autogenerated mock type for the TargetManagerInterface type
type MockTargetManager struct {
mock.Mock
}
type MockTargetManager_Expecter struct {
mock *mock.Mock
}
func (_m *MockTargetManager) EXPECT() *MockTargetManager_Expecter {
return &MockTargetManager_Expecter{mock: &_m.Mock}
}
// GetCollectionTargetVersion provides a mock function with given fields: collectionID, scope
func (_m *MockTargetManager) GetCollectionTargetVersion(collectionID int64, scope int32) int64 {
ret := _m.Called(collectionID, scope)
var r0 int64
if rf, ok := ret.Get(0).(func(int64, int32) int64); ok {
r0 = rf(collectionID, scope)
} else {
r0 = ret.Get(0).(int64)
}
return r0
}
// MockTargetManager_GetCollectionTargetVersion_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetCollectionTargetVersion'
type MockTargetManager_GetCollectionTargetVersion_Call struct {
*mock.Call
}
// GetCollectionTargetVersion is a helper method to define mock.On call
// - collectionID int64
// - scope int32
func (_e *MockTargetManager_Expecter) GetCollectionTargetVersion(collectionID interface{}, scope interface{}) *MockTargetManager_GetCollectionTargetVersion_Call {
return &MockTargetManager_GetCollectionTargetVersion_Call{Call: _e.mock.On("GetCollectionTargetVersion", collectionID, scope)}
}
func (_c *MockTargetManager_GetCollectionTargetVersion_Call) Run(run func(collectionID int64, scope int32)) *MockTargetManager_GetCollectionTargetVersion_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64), args[1].(int32))
})
return _c
}
func (_c *MockTargetManager_GetCollectionTargetVersion_Call) Return(_a0 int64) *MockTargetManager_GetCollectionTargetVersion_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockTargetManager_GetCollectionTargetVersion_Call) RunAndReturn(run func(int64, int32) int64) *MockTargetManager_GetCollectionTargetVersion_Call {
_c.Call.Return(run)
return _c
}
// GetDmChannel provides a mock function with given fields: collectionID, channel, scope
func (_m *MockTargetManager) GetDmChannel(collectionID int64, channel string, scope int32) *DmChannel {
ret := _m.Called(collectionID, channel, scope)
var r0 *DmChannel
if rf, ok := ret.Get(0).(func(int64, string, int32) *DmChannel); ok {
r0 = rf(collectionID, channel, scope)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(*DmChannel)
}
}
return r0
}
// MockTargetManager_GetDmChannel_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetDmChannel'
type MockTargetManager_GetDmChannel_Call struct {
*mock.Call
}
// GetDmChannel is a helper method to define mock.On call
// - collectionID int64
// - channel string
// - scope int32
func (_e *MockTargetManager_Expecter) GetDmChannel(collectionID interface{}, channel interface{}, scope interface{}) *MockTargetManager_GetDmChannel_Call {
return &MockTargetManager_GetDmChannel_Call{Call: _e.mock.On("GetDmChannel", collectionID, channel, scope)}
}
func (_c *MockTargetManager_GetDmChannel_Call) Run(run func(collectionID int64, channel string, scope int32)) *MockTargetManager_GetDmChannel_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64), args[1].(string), args[2].(int32))
})
return _c
}
func (_c *MockTargetManager_GetDmChannel_Call) Return(_a0 *DmChannel) *MockTargetManager_GetDmChannel_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockTargetManager_GetDmChannel_Call) RunAndReturn(run func(int64, string, int32) *DmChannel) *MockTargetManager_GetDmChannel_Call {
_c.Call.Return(run)
return _c
}
// GetDmChannelsByCollection provides a mock function with given fields: collectionID, scope
func (_m *MockTargetManager) GetDmChannelsByCollection(collectionID int64, scope int32) map[string]*DmChannel {
ret := _m.Called(collectionID, scope)
var r0 map[string]*DmChannel
if rf, ok := ret.Get(0).(func(int64, int32) map[string]*DmChannel); ok {
r0 = rf(collectionID, scope)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(map[string]*DmChannel)
}
}
return r0
}
// MockTargetManager_GetDmChannelsByCollection_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetDmChannelsByCollection'
type MockTargetManager_GetDmChannelsByCollection_Call struct {
*mock.Call
}
// GetDmChannelsByCollection is a helper method to define mock.On call
// - collectionID int64
// - scope int32
func (_e *MockTargetManager_Expecter) GetDmChannelsByCollection(collectionID interface{}, scope interface{}) *MockTargetManager_GetDmChannelsByCollection_Call {
return &MockTargetManager_GetDmChannelsByCollection_Call{Call: _e.mock.On("GetDmChannelsByCollection", collectionID, scope)}
}
func (_c *MockTargetManager_GetDmChannelsByCollection_Call) Run(run func(collectionID int64, scope int32)) *MockTargetManager_GetDmChannelsByCollection_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64), args[1].(int32))
})
return _c
}
func (_c *MockTargetManager_GetDmChannelsByCollection_Call) Return(_a0 map[string]*DmChannel) *MockTargetManager_GetDmChannelsByCollection_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockTargetManager_GetDmChannelsByCollection_Call) RunAndReturn(run func(int64, int32) map[string]*DmChannel) *MockTargetManager_GetDmChannelsByCollection_Call {
_c.Call.Return(run)
return _c
}
// GetDroppedSegmentsByChannel provides a mock function with given fields: collectionID, channelName, scope
func (_m *MockTargetManager) GetDroppedSegmentsByChannel(collectionID int64, channelName string, scope int32) []int64 {
ret := _m.Called(collectionID, channelName, scope)
var r0 []int64
if rf, ok := ret.Get(0).(func(int64, string, int32) []int64); ok {
r0 = rf(collectionID, channelName, scope)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).([]int64)
}
}
return r0
}
// MockTargetManager_GetDroppedSegmentsByChannel_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetDroppedSegmentsByChannel'
type MockTargetManager_GetDroppedSegmentsByChannel_Call struct {
*mock.Call
}
// GetDroppedSegmentsByChannel is a helper method to define mock.On call
// - collectionID int64
// - channelName string
// - scope int32
func (_e *MockTargetManager_Expecter) GetDroppedSegmentsByChannel(collectionID interface{}, channelName interface{}, scope interface{}) *MockTargetManager_GetDroppedSegmentsByChannel_Call {
return &MockTargetManager_GetDroppedSegmentsByChannel_Call{Call: _e.mock.On("GetDroppedSegmentsByChannel", collectionID, channelName, scope)}
}
func (_c *MockTargetManager_GetDroppedSegmentsByChannel_Call) Run(run func(collectionID int64, channelName string, scope int32)) *MockTargetManager_GetDroppedSegmentsByChannel_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64), args[1].(string), args[2].(int32))
})
return _c
}
func (_c *MockTargetManager_GetDroppedSegmentsByChannel_Call) Return(_a0 []int64) *MockTargetManager_GetDroppedSegmentsByChannel_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockTargetManager_GetDroppedSegmentsByChannel_Call) RunAndReturn(run func(int64, string, int32) []int64) *MockTargetManager_GetDroppedSegmentsByChannel_Call {
_c.Call.Return(run)
return _c
}
// GetGrowingSegmentsByChannel provides a mock function with given fields: collectionID, channelName, scope
func (_m *MockTargetManager) GetGrowingSegmentsByChannel(collectionID int64, channelName string, scope int32) typeutil.Set[int64] {
ret := _m.Called(collectionID, channelName, scope)
var r0 typeutil.Set[int64]
if rf, ok := ret.Get(0).(func(int64, string, int32) typeutil.Set[int64]); ok {
r0 = rf(collectionID, channelName, scope)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(typeutil.Set[int64])
}
}
return r0
}
// MockTargetManager_GetGrowingSegmentsByChannel_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetGrowingSegmentsByChannel'
type MockTargetManager_GetGrowingSegmentsByChannel_Call struct {
*mock.Call
}
// GetGrowingSegmentsByChannel is a helper method to define mock.On call
// - collectionID int64
// - channelName string
// - scope int32
func (_e *MockTargetManager_Expecter) GetGrowingSegmentsByChannel(collectionID interface{}, channelName interface{}, scope interface{}) *MockTargetManager_GetGrowingSegmentsByChannel_Call {
return &MockTargetManager_GetGrowingSegmentsByChannel_Call{Call: _e.mock.On("GetGrowingSegmentsByChannel", collectionID, channelName, scope)}
}
func (_c *MockTargetManager_GetGrowingSegmentsByChannel_Call) Run(run func(collectionID int64, channelName string, scope int32)) *MockTargetManager_GetGrowingSegmentsByChannel_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64), args[1].(string), args[2].(int32))
})
return _c
}
func (_c *MockTargetManager_GetGrowingSegmentsByChannel_Call) Return(_a0 typeutil.Set[int64]) *MockTargetManager_GetGrowingSegmentsByChannel_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockTargetManager_GetGrowingSegmentsByChannel_Call) RunAndReturn(run func(int64, string, int32) typeutil.Set[int64]) *MockTargetManager_GetGrowingSegmentsByChannel_Call {
_c.Call.Return(run)
return _c
}
// GetGrowingSegmentsByCollection provides a mock function with given fields: collectionID, scope
func (_m *MockTargetManager) GetGrowingSegmentsByCollection(collectionID int64, scope int32) typeutil.Set[int64] {
ret := _m.Called(collectionID, scope)
var r0 typeutil.Set[int64]
if rf, ok := ret.Get(0).(func(int64, int32) typeutil.Set[int64]); ok {
r0 = rf(collectionID, scope)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(typeutil.Set[int64])
}
}
return r0
}
// MockTargetManager_GetGrowingSegmentsByCollection_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetGrowingSegmentsByCollection'
type MockTargetManager_GetGrowingSegmentsByCollection_Call struct {
*mock.Call
}
// GetGrowingSegmentsByCollection is a helper method to define mock.On call
// - collectionID int64
// - scope int32
func (_e *MockTargetManager_Expecter) GetGrowingSegmentsByCollection(collectionID interface{}, scope interface{}) *MockTargetManager_GetGrowingSegmentsByCollection_Call {
return &MockTargetManager_GetGrowingSegmentsByCollection_Call{Call: _e.mock.On("GetGrowingSegmentsByCollection", collectionID, scope)}
}
func (_c *MockTargetManager_GetGrowingSegmentsByCollection_Call) Run(run func(collectionID int64, scope int32)) *MockTargetManager_GetGrowingSegmentsByCollection_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64), args[1].(int32))
})
return _c
}
func (_c *MockTargetManager_GetGrowingSegmentsByCollection_Call) Return(_a0 typeutil.Set[int64]) *MockTargetManager_GetGrowingSegmentsByCollection_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockTargetManager_GetGrowingSegmentsByCollection_Call) RunAndReturn(run func(int64, int32) typeutil.Set[int64]) *MockTargetManager_GetGrowingSegmentsByCollection_Call {
_c.Call.Return(run)
return _c
}
// GetSealedSegment provides a mock function with given fields: collectionID, id, scope
func (_m *MockTargetManager) GetSealedSegment(collectionID int64, id int64, scope int32) *datapb.SegmentInfo {
ret := _m.Called(collectionID, id, scope)
var r0 *datapb.SegmentInfo
if rf, ok := ret.Get(0).(func(int64, int64, int32) *datapb.SegmentInfo); ok {
r0 = rf(collectionID, id, scope)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(*datapb.SegmentInfo)
}
}
return r0
}
// MockTargetManager_GetSealedSegment_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetSealedSegment'
type MockTargetManager_GetSealedSegment_Call struct {
*mock.Call
}
// GetSealedSegment is a helper method to define mock.On call
// - collectionID int64
// - id int64
// - scope int32
func (_e *MockTargetManager_Expecter) GetSealedSegment(collectionID interface{}, id interface{}, scope interface{}) *MockTargetManager_GetSealedSegment_Call {
return &MockTargetManager_GetSealedSegment_Call{Call: _e.mock.On("GetSealedSegment", collectionID, id, scope)}
}
func (_c *MockTargetManager_GetSealedSegment_Call) Run(run func(collectionID int64, id int64, scope int32)) *MockTargetManager_GetSealedSegment_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64), args[1].(int64), args[2].(int32))
})
return _c
}
func (_c *MockTargetManager_GetSealedSegment_Call) Return(_a0 *datapb.SegmentInfo) *MockTargetManager_GetSealedSegment_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockTargetManager_GetSealedSegment_Call) RunAndReturn(run func(int64, int64, int32) *datapb.SegmentInfo) *MockTargetManager_GetSealedSegment_Call {
_c.Call.Return(run)
return _c
}
// GetSealedSegmentsByChannel provides a mock function with given fields: collectionID, channelName, scope
func (_m *MockTargetManager) GetSealedSegmentsByChannel(collectionID int64, channelName string, scope int32) map[int64]*datapb.SegmentInfo {
ret := _m.Called(collectionID, channelName, scope)
var r0 map[int64]*datapb.SegmentInfo
if rf, ok := ret.Get(0).(func(int64, string, int32) map[int64]*datapb.SegmentInfo); ok {
r0 = rf(collectionID, channelName, scope)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(map[int64]*datapb.SegmentInfo)
}
}
return r0
}
// MockTargetManager_GetSealedSegmentsByChannel_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetSealedSegmentsByChannel'
type MockTargetManager_GetSealedSegmentsByChannel_Call struct {
*mock.Call
}
// GetSealedSegmentsByChannel is a helper method to define mock.On call
// - collectionID int64
// - channelName string
// - scope int32
func (_e *MockTargetManager_Expecter) GetSealedSegmentsByChannel(collectionID interface{}, channelName interface{}, scope interface{}) *MockTargetManager_GetSealedSegmentsByChannel_Call {
return &MockTargetManager_GetSealedSegmentsByChannel_Call{Call: _e.mock.On("GetSealedSegmentsByChannel", collectionID, channelName, scope)}
}
func (_c *MockTargetManager_GetSealedSegmentsByChannel_Call) Run(run func(collectionID int64, channelName string, scope int32)) *MockTargetManager_GetSealedSegmentsByChannel_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64), args[1].(string), args[2].(int32))
})
return _c
}
func (_c *MockTargetManager_GetSealedSegmentsByChannel_Call) Return(_a0 map[int64]*datapb.SegmentInfo) *MockTargetManager_GetSealedSegmentsByChannel_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockTargetManager_GetSealedSegmentsByChannel_Call) RunAndReturn(run func(int64, string, int32) map[int64]*datapb.SegmentInfo) *MockTargetManager_GetSealedSegmentsByChannel_Call {
_c.Call.Return(run)
return _c
}
// GetSealedSegmentsByCollection provides a mock function with given fields: collectionID, scope
func (_m *MockTargetManager) GetSealedSegmentsByCollection(collectionID int64, scope int32) map[int64]*datapb.SegmentInfo {
ret := _m.Called(collectionID, scope)
var r0 map[int64]*datapb.SegmentInfo
if rf, ok := ret.Get(0).(func(int64, int32) map[int64]*datapb.SegmentInfo); ok {
r0 = rf(collectionID, scope)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(map[int64]*datapb.SegmentInfo)
}
}
return r0
}
// MockTargetManager_GetSealedSegmentsByCollection_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetSealedSegmentsByCollection'
type MockTargetManager_GetSealedSegmentsByCollection_Call struct {
*mock.Call
}
// GetSealedSegmentsByCollection is a helper method to define mock.On call
// - collectionID int64
// - scope int32
func (_e *MockTargetManager_Expecter) GetSealedSegmentsByCollection(collectionID interface{}, scope interface{}) *MockTargetManager_GetSealedSegmentsByCollection_Call {
return &MockTargetManager_GetSealedSegmentsByCollection_Call{Call: _e.mock.On("GetSealedSegmentsByCollection", collectionID, scope)}
}
func (_c *MockTargetManager_GetSealedSegmentsByCollection_Call) Run(run func(collectionID int64, scope int32)) *MockTargetManager_GetSealedSegmentsByCollection_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64), args[1].(int32))
})
return _c
}
func (_c *MockTargetManager_GetSealedSegmentsByCollection_Call) Return(_a0 map[int64]*datapb.SegmentInfo) *MockTargetManager_GetSealedSegmentsByCollection_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockTargetManager_GetSealedSegmentsByCollection_Call) RunAndReturn(run func(int64, int32) map[int64]*datapb.SegmentInfo) *MockTargetManager_GetSealedSegmentsByCollection_Call {
_c.Call.Return(run)
return _c
}
// GetSealedSegmentsByPartition provides a mock function with given fields: collectionID, partitionID, scope
func (_m *MockTargetManager) GetSealedSegmentsByPartition(collectionID int64, partitionID int64, scope int32) map[int64]*datapb.SegmentInfo {
ret := _m.Called(collectionID, partitionID, scope)
var r0 map[int64]*datapb.SegmentInfo
if rf, ok := ret.Get(0).(func(int64, int64, int32) map[int64]*datapb.SegmentInfo); ok {
r0 = rf(collectionID, partitionID, scope)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(map[int64]*datapb.SegmentInfo)
}
}
return r0
}
// MockTargetManager_GetSealedSegmentsByPartition_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetSealedSegmentsByPartition'
type MockTargetManager_GetSealedSegmentsByPartition_Call struct {
*mock.Call
}
// GetSealedSegmentsByPartition is a helper method to define mock.On call
// - collectionID int64
// - partitionID int64
// - scope int32
func (_e *MockTargetManager_Expecter) GetSealedSegmentsByPartition(collectionID interface{}, partitionID interface{}, scope interface{}) *MockTargetManager_GetSealedSegmentsByPartition_Call {
return &MockTargetManager_GetSealedSegmentsByPartition_Call{Call: _e.mock.On("GetSealedSegmentsByPartition", collectionID, partitionID, scope)}
}
func (_c *MockTargetManager_GetSealedSegmentsByPartition_Call) Run(run func(collectionID int64, partitionID int64, scope int32)) *MockTargetManager_GetSealedSegmentsByPartition_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64), args[1].(int64), args[2].(int32))
})
return _c
}
func (_c *MockTargetManager_GetSealedSegmentsByPartition_Call) Return(_a0 map[int64]*datapb.SegmentInfo) *MockTargetManager_GetSealedSegmentsByPartition_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockTargetManager_GetSealedSegmentsByPartition_Call) RunAndReturn(run func(int64, int64, int32) map[int64]*datapb.SegmentInfo) *MockTargetManager_GetSealedSegmentsByPartition_Call {
_c.Call.Return(run)
return _c
}
// IsCurrentTargetExist provides a mock function with given fields: collectionID, partitionID
func (_m *MockTargetManager) IsCurrentTargetExist(collectionID int64, partitionID int64) bool {
ret := _m.Called(collectionID, partitionID)
var r0 bool
if rf, ok := ret.Get(0).(func(int64, int64) bool); ok {
r0 = rf(collectionID, partitionID)
} else {
r0 = ret.Get(0).(bool)
}
return r0
}
// MockTargetManager_IsCurrentTargetExist_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'IsCurrentTargetExist'
type MockTargetManager_IsCurrentTargetExist_Call struct {
*mock.Call
}
// IsCurrentTargetExist is a helper method to define mock.On call
// - collectionID int64
// - partitionID int64
func (_e *MockTargetManager_Expecter) IsCurrentTargetExist(collectionID interface{}, partitionID interface{}) *MockTargetManager_IsCurrentTargetExist_Call {
return &MockTargetManager_IsCurrentTargetExist_Call{Call: _e.mock.On("IsCurrentTargetExist", collectionID, partitionID)}
}
func (_c *MockTargetManager_IsCurrentTargetExist_Call) Run(run func(collectionID int64, partitionID int64)) *MockTargetManager_IsCurrentTargetExist_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64), args[1].(int64))
})
return _c
}
func (_c *MockTargetManager_IsCurrentTargetExist_Call) Return(_a0 bool) *MockTargetManager_IsCurrentTargetExist_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockTargetManager_IsCurrentTargetExist_Call) RunAndReturn(run func(int64, int64) bool) *MockTargetManager_IsCurrentTargetExist_Call {
_c.Call.Return(run)
return _c
}
// IsNextTargetExist provides a mock function with given fields: collectionID
func (_m *MockTargetManager) IsNextTargetExist(collectionID int64) bool {
ret := _m.Called(collectionID)
var r0 bool
if rf, ok := ret.Get(0).(func(int64) bool); ok {
r0 = rf(collectionID)
} else {
r0 = ret.Get(0).(bool)
}
return r0
}
// MockTargetManager_IsNextTargetExist_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'IsNextTargetExist'
type MockTargetManager_IsNextTargetExist_Call struct {
*mock.Call
}
// IsNextTargetExist is a helper method to define mock.On call
// - collectionID int64
func (_e *MockTargetManager_Expecter) IsNextTargetExist(collectionID interface{}) *MockTargetManager_IsNextTargetExist_Call {
return &MockTargetManager_IsNextTargetExist_Call{Call: _e.mock.On("IsNextTargetExist", collectionID)}
}
func (_c *MockTargetManager_IsNextTargetExist_Call) Run(run func(collectionID int64)) *MockTargetManager_IsNextTargetExist_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64))
})
return _c
}
func (_c *MockTargetManager_IsNextTargetExist_Call) Return(_a0 bool) *MockTargetManager_IsNextTargetExist_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockTargetManager_IsNextTargetExist_Call) RunAndReturn(run func(int64) bool) *MockTargetManager_IsNextTargetExist_Call {
_c.Call.Return(run)
return _c
}
// PullNextTargetV1 provides a mock function with given fields: broker, collectionID, chosenPartitionIDs
func (_m *MockTargetManager) PullNextTargetV1(broker Broker, collectionID int64, chosenPartitionIDs ...int64) (map[int64]*datapb.SegmentInfo, map[string]*DmChannel, error) {
_va := make([]interface{}, len(chosenPartitionIDs))
for _i := range chosenPartitionIDs {
_va[_i] = chosenPartitionIDs[_i]
}
var _ca []interface{}
_ca = append(_ca, broker, collectionID)
_ca = append(_ca, _va...)
ret := _m.Called(_ca...)
var r0 map[int64]*datapb.SegmentInfo
var r1 map[string]*DmChannel
var r2 error
if rf, ok := ret.Get(0).(func(Broker, int64, ...int64) (map[int64]*datapb.SegmentInfo, map[string]*DmChannel, error)); ok {
return rf(broker, collectionID, chosenPartitionIDs...)
}
if rf, ok := ret.Get(0).(func(Broker, int64, ...int64) map[int64]*datapb.SegmentInfo); ok {
r0 = rf(broker, collectionID, chosenPartitionIDs...)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(map[int64]*datapb.SegmentInfo)
}
}
if rf, ok := ret.Get(1).(func(Broker, int64, ...int64) map[string]*DmChannel); ok {
r1 = rf(broker, collectionID, chosenPartitionIDs...)
} else {
if ret.Get(1) != nil {
r1 = ret.Get(1).(map[string]*DmChannel)
}
}
if rf, ok := ret.Get(2).(func(Broker, int64, ...int64) error); ok {
r2 = rf(broker, collectionID, chosenPartitionIDs...)
} else {
r2 = ret.Error(2)
}
return r0, r1, r2
}
// MockTargetManager_PullNextTargetV1_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'PullNextTargetV1'
type MockTargetManager_PullNextTargetV1_Call struct {
*mock.Call
}
// PullNextTargetV1 is a helper method to define mock.On call
// - broker Broker
// - collectionID int64
// - chosenPartitionIDs ...int64
func (_e *MockTargetManager_Expecter) PullNextTargetV1(broker interface{}, collectionID interface{}, chosenPartitionIDs ...interface{}) *MockTargetManager_PullNextTargetV1_Call {
return &MockTargetManager_PullNextTargetV1_Call{Call: _e.mock.On("PullNextTargetV1",
append([]interface{}{broker, collectionID}, chosenPartitionIDs...)...)}
}
func (_c *MockTargetManager_PullNextTargetV1_Call) Run(run func(broker Broker, collectionID int64, chosenPartitionIDs ...int64)) *MockTargetManager_PullNextTargetV1_Call {
_c.Call.Run(func(args mock.Arguments) {
variadicArgs := make([]int64, len(args)-2)
for i, a := range args[2:] {
if a != nil {
variadicArgs[i] = a.(int64)
}
}
run(args[0].(Broker), args[1].(int64), variadicArgs...)
})
return _c
}
func (_c *MockTargetManager_PullNextTargetV1_Call) Return(_a0 map[int64]*datapb.SegmentInfo, _a1 map[string]*DmChannel, _a2 error) *MockTargetManager_PullNextTargetV1_Call {
_c.Call.Return(_a0, _a1, _a2)
return _c
}
func (_c *MockTargetManager_PullNextTargetV1_Call) RunAndReturn(run func(Broker, int64, ...int64) (map[int64]*datapb.SegmentInfo, map[string]*DmChannel, error)) *MockTargetManager_PullNextTargetV1_Call {
_c.Call.Return(run)
return _c
}
// PullNextTargetV2 provides a mock function with given fields: broker, collectionID, chosenPartitionIDs
func (_m *MockTargetManager) PullNextTargetV2(broker Broker, collectionID int64, chosenPartitionIDs ...int64) (map[int64]*datapb.SegmentInfo, map[string]*DmChannel, error) {
_va := make([]interface{}, len(chosenPartitionIDs))
for _i := range chosenPartitionIDs {
_va[_i] = chosenPartitionIDs[_i]
}
var _ca []interface{}
_ca = append(_ca, broker, collectionID)
_ca = append(_ca, _va...)
ret := _m.Called(_ca...)
var r0 map[int64]*datapb.SegmentInfo
var r1 map[string]*DmChannel
var r2 error
if rf, ok := ret.Get(0).(func(Broker, int64, ...int64) (map[int64]*datapb.SegmentInfo, map[string]*DmChannel, error)); ok {
return rf(broker, collectionID, chosenPartitionIDs...)
}
if rf, ok := ret.Get(0).(func(Broker, int64, ...int64) map[int64]*datapb.SegmentInfo); ok {
r0 = rf(broker, collectionID, chosenPartitionIDs...)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(map[int64]*datapb.SegmentInfo)
}
}
if rf, ok := ret.Get(1).(func(Broker, int64, ...int64) map[string]*DmChannel); ok {
r1 = rf(broker, collectionID, chosenPartitionIDs...)
} else {
if ret.Get(1) != nil {
r1 = ret.Get(1).(map[string]*DmChannel)
}
}
if rf, ok := ret.Get(2).(func(Broker, int64, ...int64) error); ok {
r2 = rf(broker, collectionID, chosenPartitionIDs...)
} else {
r2 = ret.Error(2)
}
return r0, r1, r2
}
// MockTargetManager_PullNextTargetV2_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'PullNextTargetV2'
type MockTargetManager_PullNextTargetV2_Call struct {
*mock.Call
}
// PullNextTargetV2 is a helper method to define mock.On call
// - broker Broker
// - collectionID int64
// - chosenPartitionIDs ...int64
func (_e *MockTargetManager_Expecter) PullNextTargetV2(broker interface{}, collectionID interface{}, chosenPartitionIDs ...interface{}) *MockTargetManager_PullNextTargetV2_Call {
return &MockTargetManager_PullNextTargetV2_Call{Call: _e.mock.On("PullNextTargetV2",
append([]interface{}{broker, collectionID}, chosenPartitionIDs...)...)}
}
func (_c *MockTargetManager_PullNextTargetV2_Call) Run(run func(broker Broker, collectionID int64, chosenPartitionIDs ...int64)) *MockTargetManager_PullNextTargetV2_Call {
_c.Call.Run(func(args mock.Arguments) {
variadicArgs := make([]int64, len(args)-2)
for i, a := range args[2:] {
if a != nil {
variadicArgs[i] = a.(int64)
}
}
run(args[0].(Broker), args[1].(int64), variadicArgs...)
})
return _c
}
func (_c *MockTargetManager_PullNextTargetV2_Call) Return(_a0 map[int64]*datapb.SegmentInfo, _a1 map[string]*DmChannel, _a2 error) *MockTargetManager_PullNextTargetV2_Call {
_c.Call.Return(_a0, _a1, _a2)
return _c
}
func (_c *MockTargetManager_PullNextTargetV2_Call) RunAndReturn(run func(Broker, int64, ...int64) (map[int64]*datapb.SegmentInfo, map[string]*DmChannel, error)) *MockTargetManager_PullNextTargetV2_Call {
_c.Call.Return(run)
return _c
}
// Recover provides a mock function with given fields: catalog
func (_m *MockTargetManager) Recover(catalog metastore.QueryCoordCatalog) error {
ret := _m.Called(catalog)
var r0 error
if rf, ok := ret.Get(0).(func(metastore.QueryCoordCatalog) error); ok {
r0 = rf(catalog)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockTargetManager_Recover_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Recover'
type MockTargetManager_Recover_Call struct {
*mock.Call
}
// Recover is a helper method to define mock.On call
// - catalog metastore.QueryCoordCatalog
func (_e *MockTargetManager_Expecter) Recover(catalog interface{}) *MockTargetManager_Recover_Call {
return &MockTargetManager_Recover_Call{Call: _e.mock.On("Recover", catalog)}
}
func (_c *MockTargetManager_Recover_Call) Run(run func(catalog metastore.QueryCoordCatalog)) *MockTargetManager_Recover_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(metastore.QueryCoordCatalog))
})
return _c
}
func (_c *MockTargetManager_Recover_Call) Return(_a0 error) *MockTargetManager_Recover_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockTargetManager_Recover_Call) RunAndReturn(run func(metastore.QueryCoordCatalog) error) *MockTargetManager_Recover_Call {
_c.Call.Return(run)
return _c
}
// RemoveCollection provides a mock function with given fields: collectionID
func (_m *MockTargetManager) RemoveCollection(collectionID int64) {
_m.Called(collectionID)
}
// MockTargetManager_RemoveCollection_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RemoveCollection'
type MockTargetManager_RemoveCollection_Call struct {
*mock.Call
}
// RemoveCollection is a helper method to define mock.On call
// - collectionID int64
func (_e *MockTargetManager_Expecter) RemoveCollection(collectionID interface{}) *MockTargetManager_RemoveCollection_Call {
return &MockTargetManager_RemoveCollection_Call{Call: _e.mock.On("RemoveCollection", collectionID)}
}
func (_c *MockTargetManager_RemoveCollection_Call) Run(run func(collectionID int64)) *MockTargetManager_RemoveCollection_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64))
})
return _c
}
func (_c *MockTargetManager_RemoveCollection_Call) Return() *MockTargetManager_RemoveCollection_Call {
_c.Call.Return()
return _c
}
func (_c *MockTargetManager_RemoveCollection_Call) RunAndReturn(run func(int64)) *MockTargetManager_RemoveCollection_Call {
_c.Call.Return(run)
return _c
}
// RemovePartition provides a mock function with given fields: collectionID, partitionIDs
func (_m *MockTargetManager) RemovePartition(collectionID int64, partitionIDs ...int64) {
_va := make([]interface{}, len(partitionIDs))
for _i := range partitionIDs {
_va[_i] = partitionIDs[_i]
}
var _ca []interface{}
_ca = append(_ca, collectionID)
_ca = append(_ca, _va...)
_m.Called(_ca...)
}
// MockTargetManager_RemovePartition_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RemovePartition'
type MockTargetManager_RemovePartition_Call struct {
*mock.Call
}
// RemovePartition is a helper method to define mock.On call
// - collectionID int64
// - partitionIDs ...int64
func (_e *MockTargetManager_Expecter) RemovePartition(collectionID interface{}, partitionIDs ...interface{}) *MockTargetManager_RemovePartition_Call {
return &MockTargetManager_RemovePartition_Call{Call: _e.mock.On("RemovePartition",
append([]interface{}{collectionID}, partitionIDs...)...)}
}
func (_c *MockTargetManager_RemovePartition_Call) Run(run func(collectionID int64, partitionIDs ...int64)) *MockTargetManager_RemovePartition_Call {
_c.Call.Run(func(args mock.Arguments) {
variadicArgs := make([]int64, len(args)-1)
for i, a := range args[1:] {
if a != nil {
variadicArgs[i] = a.(int64)
}
}
run(args[0].(int64), variadicArgs...)
})
return _c
}
func (_c *MockTargetManager_RemovePartition_Call) Return() *MockTargetManager_RemovePartition_Call {
_c.Call.Return()
return _c
}
func (_c *MockTargetManager_RemovePartition_Call) RunAndReturn(run func(int64, ...int64)) *MockTargetManager_RemovePartition_Call {
_c.Call.Return(run)
return _c
}
// SaveCurrentTarget provides a mock function with given fields: catalog
func (_m *MockTargetManager) SaveCurrentTarget(catalog metastore.QueryCoordCatalog) {
_m.Called(catalog)
}
// MockTargetManager_SaveCurrentTarget_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SaveCurrentTarget'
type MockTargetManager_SaveCurrentTarget_Call struct {
*mock.Call
}
// SaveCurrentTarget is a helper method to define mock.On call
// - catalog metastore.QueryCoordCatalog
func (_e *MockTargetManager_Expecter) SaveCurrentTarget(catalog interface{}) *MockTargetManager_SaveCurrentTarget_Call {
return &MockTargetManager_SaveCurrentTarget_Call{Call: _e.mock.On("SaveCurrentTarget", catalog)}
}
func (_c *MockTargetManager_SaveCurrentTarget_Call) Run(run func(catalog metastore.QueryCoordCatalog)) *MockTargetManager_SaveCurrentTarget_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(metastore.QueryCoordCatalog))
})
return _c
}
func (_c *MockTargetManager_SaveCurrentTarget_Call) Return() *MockTargetManager_SaveCurrentTarget_Call {
_c.Call.Return()
return _c
}
func (_c *MockTargetManager_SaveCurrentTarget_Call) RunAndReturn(run func(metastore.QueryCoordCatalog)) *MockTargetManager_SaveCurrentTarget_Call {
_c.Call.Return(run)
return _c
}
// UpdateCollectionCurrentTarget provides a mock function with given fields: collectionID
func (_m *MockTargetManager) UpdateCollectionCurrentTarget(collectionID int64) bool {
ret := _m.Called(collectionID)
var r0 bool
if rf, ok := ret.Get(0).(func(int64) bool); ok {
r0 = rf(collectionID)
} else {
r0 = ret.Get(0).(bool)
}
return r0
}
// MockTargetManager_UpdateCollectionCurrentTarget_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'UpdateCollectionCurrentTarget'
type MockTargetManager_UpdateCollectionCurrentTarget_Call struct {
*mock.Call
}
// UpdateCollectionCurrentTarget is a helper method to define mock.On call
// - collectionID int64
func (_e *MockTargetManager_Expecter) UpdateCollectionCurrentTarget(collectionID interface{}) *MockTargetManager_UpdateCollectionCurrentTarget_Call {
return &MockTargetManager_UpdateCollectionCurrentTarget_Call{Call: _e.mock.On("UpdateCollectionCurrentTarget", collectionID)}
}
func (_c *MockTargetManager_UpdateCollectionCurrentTarget_Call) Run(run func(collectionID int64)) *MockTargetManager_UpdateCollectionCurrentTarget_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64))
})
return _c
}
func (_c *MockTargetManager_UpdateCollectionCurrentTarget_Call) Return(_a0 bool) *MockTargetManager_UpdateCollectionCurrentTarget_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockTargetManager_UpdateCollectionCurrentTarget_Call) RunAndReturn(run func(int64) bool) *MockTargetManager_UpdateCollectionCurrentTarget_Call {
_c.Call.Return(run)
return _c
}
// UpdateCollectionNextTarget provides a mock function with given fields: collectionID
func (_m *MockTargetManager) UpdateCollectionNextTarget(collectionID int64) error {
ret := _m.Called(collectionID)
var r0 error
if rf, ok := ret.Get(0).(func(int64) error); ok {
r0 = rf(collectionID)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockTargetManager_UpdateCollectionNextTarget_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'UpdateCollectionNextTarget'
type MockTargetManager_UpdateCollectionNextTarget_Call struct {
*mock.Call
}
// UpdateCollectionNextTarget is a helper method to define mock.On call
// - collectionID int64
func (_e *MockTargetManager_Expecter) UpdateCollectionNextTarget(collectionID interface{}) *MockTargetManager_UpdateCollectionNextTarget_Call {
return &MockTargetManager_UpdateCollectionNextTarget_Call{Call: _e.mock.On("UpdateCollectionNextTarget", collectionID)}
}
func (_c *MockTargetManager_UpdateCollectionNextTarget_Call) Run(run func(collectionID int64)) *MockTargetManager_UpdateCollectionNextTarget_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64))
})
return _c
}
func (_c *MockTargetManager_UpdateCollectionNextTarget_Call) Return(_a0 error) *MockTargetManager_UpdateCollectionNextTarget_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockTargetManager_UpdateCollectionNextTarget_Call) RunAndReturn(run func(int64) error) *MockTargetManager_UpdateCollectionNextTarget_Call {
_c.Call.Return(run)
return _c
}
// NewMockTargetManager creates a new instance of MockTargetManager. 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 NewMockTargetManager(t interface {
mock.TestingT
Cleanup(func())
}) *MockTargetManager {
mock := &MockTargetManager{}
mock.Mock.Test(t)
t.Cleanup(func() { mock.AssertExpectations(t) })
return mock
}