milvus/internal/datacoord/session/mock_cluster.go
Bingyi Sun f9827392bb
enhance: implement external collection update task with source change detection (#45905)
issue: #45881 
Add persistent task management for external collections with automatic
detection of external_source and external_spec changes. When source
changes, the system aborts running tasks and creates new ones, ensuring
only one active task per collection. Tasks validate their source on
completion to prevent superseded tasks from committing results.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
- Core invariant: at most one active UpdateExternalCollection task
exists per collection — tasks are serialized by collectionID
(collection-level locking) and any change to external_source or
external_spec aborts superseded tasks and causes a new task creation
(externalCollectionManager + external_collection_task_meta
collection-based locks enforce this).
- What was simplified/removed: per-task fine-grained locking and
concurrent multi-task acceptance per collection were replaced by
collection-level synchronization (external_collection_task_meta.go) and
a single persistent task lifecycle in DataCoord/Index task code;
redundant double-concurrent update paths were removed by checking
existing task presence in AddTask/LoadOrStore and aborting/overwriting
via Drop/Cancel flows.
- Why this does NOT cause data loss or regress behavior: task state
transitions and commit are validated against the current external
source/spec before applying changes — UpdateStateWithMeta and SetJobInfo
verify task metadata and persist via catalog only under matching
collection-state; DataNode externalCollectionManager persists task
results to in-memory manager and exposes Query/Drop flows (services.go)
without modifying existing segment data unless a task successfully
finishes and SetJobInfo atomically updates segments via meta/catalog
calls, preventing superseded tasks from committing stale results.
- New capability added: end-to-end external collection update workflow —
DataCoord Index task + Cluster RPC helpers + DataNode external task
runner and ExternalCollectionManager enable creating, querying,
cancelling, and applying external collection updates
(fragment-to-segment balancing, kept/updated segment handling, allocator
integration); accompanying unit tests cover success, failure,
cancellation, allocator errors, and balancing logic.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: sunby <sunbingyi1992@gmail.com>
2025-12-29 19:53:21 +08:00

1111 lines
35 KiB
Go

// Code generated by mockery v2.53.3. DO NOT EDIT.
package session
import (
datapb "github.com/milvus-io/milvus/pkg/v2/proto/datapb"
mock "github.com/stretchr/testify/mock"
workerpb "github.com/milvus-io/milvus/pkg/v2/proto/workerpb"
)
// MockCluster is an autogenerated mock type for the Cluster type
type MockCluster struct {
mock.Mock
}
type MockCluster_Expecter struct {
mock *mock.Mock
}
func (_m *MockCluster) EXPECT() *MockCluster_Expecter {
return &MockCluster_Expecter{mock: &_m.Mock}
}
// CreateAnalyze provides a mock function with given fields: nodeID, in
func (_m *MockCluster) CreateAnalyze(nodeID int64, in *workerpb.AnalyzeRequest) error {
ret := _m.Called(nodeID, in)
if len(ret) == 0 {
panic("no return value specified for CreateAnalyze")
}
var r0 error
if rf, ok := ret.Get(0).(func(int64, *workerpb.AnalyzeRequest) error); ok {
r0 = rf(nodeID, in)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockCluster_CreateAnalyze_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CreateAnalyze'
type MockCluster_CreateAnalyze_Call struct {
*mock.Call
}
// CreateAnalyze is a helper method to define mock.On call
// - nodeID int64
// - in *workerpb.AnalyzeRequest
func (_e *MockCluster_Expecter) CreateAnalyze(nodeID interface{}, in interface{}) *MockCluster_CreateAnalyze_Call {
return &MockCluster_CreateAnalyze_Call{Call: _e.mock.On("CreateAnalyze", nodeID, in)}
}
func (_c *MockCluster_CreateAnalyze_Call) Run(run func(nodeID int64, in *workerpb.AnalyzeRequest)) *MockCluster_CreateAnalyze_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64), args[1].(*workerpb.AnalyzeRequest))
})
return _c
}
func (_c *MockCluster_CreateAnalyze_Call) Return(_a0 error) *MockCluster_CreateAnalyze_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockCluster_CreateAnalyze_Call) RunAndReturn(run func(int64, *workerpb.AnalyzeRequest) error) *MockCluster_CreateAnalyze_Call {
_c.Call.Return(run)
return _c
}
// CreateCompaction provides a mock function with given fields: nodeID, in
func (_m *MockCluster) CreateCompaction(nodeID int64, in *datapb.CompactionPlan) error {
ret := _m.Called(nodeID, in)
if len(ret) == 0 {
panic("no return value specified for CreateCompaction")
}
var r0 error
if rf, ok := ret.Get(0).(func(int64, *datapb.CompactionPlan) error); ok {
r0 = rf(nodeID, in)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockCluster_CreateCompaction_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CreateCompaction'
type MockCluster_CreateCompaction_Call struct {
*mock.Call
}
// CreateCompaction is a helper method to define mock.On call
// - nodeID int64
// - in *datapb.CompactionPlan
func (_e *MockCluster_Expecter) CreateCompaction(nodeID interface{}, in interface{}) *MockCluster_CreateCompaction_Call {
return &MockCluster_CreateCompaction_Call{Call: _e.mock.On("CreateCompaction", nodeID, in)}
}
func (_c *MockCluster_CreateCompaction_Call) Run(run func(nodeID int64, in *datapb.CompactionPlan)) *MockCluster_CreateCompaction_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64), args[1].(*datapb.CompactionPlan))
})
return _c
}
func (_c *MockCluster_CreateCompaction_Call) Return(_a0 error) *MockCluster_CreateCompaction_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockCluster_CreateCompaction_Call) RunAndReturn(run func(int64, *datapb.CompactionPlan) error) *MockCluster_CreateCompaction_Call {
_c.Call.Return(run)
return _c
}
// CreateExternalCollectionTask provides a mock function with given fields: nodeID, req
func (_m *MockCluster) CreateExternalCollectionTask(nodeID int64, req *datapb.UpdateExternalCollectionRequest) error {
ret := _m.Called(nodeID, req)
if len(ret) == 0 {
panic("no return value specified for CreateExternalCollectionTask")
}
var r0 error
if rf, ok := ret.Get(0).(func(int64, *datapb.UpdateExternalCollectionRequest) error); ok {
r0 = rf(nodeID, req)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockCluster_CreateExternalCollectionTask_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CreateExternalCollectionTask'
type MockCluster_CreateExternalCollectionTask_Call struct {
*mock.Call
}
// CreateExternalCollectionTask is a helper method to define mock.On call
// - nodeID int64
// - req *datapb.UpdateExternalCollectionRequest
func (_e *MockCluster_Expecter) CreateExternalCollectionTask(nodeID interface{}, req interface{}) *MockCluster_CreateExternalCollectionTask_Call {
return &MockCluster_CreateExternalCollectionTask_Call{Call: _e.mock.On("CreateExternalCollectionTask", nodeID, req)}
}
func (_c *MockCluster_CreateExternalCollectionTask_Call) Run(run func(nodeID int64, req *datapb.UpdateExternalCollectionRequest)) *MockCluster_CreateExternalCollectionTask_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64), args[1].(*datapb.UpdateExternalCollectionRequest))
})
return _c
}
func (_c *MockCluster_CreateExternalCollectionTask_Call) Return(_a0 error) *MockCluster_CreateExternalCollectionTask_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockCluster_CreateExternalCollectionTask_Call) RunAndReturn(run func(int64, *datapb.UpdateExternalCollectionRequest) error) *MockCluster_CreateExternalCollectionTask_Call {
_c.Call.Return(run)
return _c
}
// CreateImport provides a mock function with given fields: nodeID, in, taskSlot
func (_m *MockCluster) CreateImport(nodeID int64, in *datapb.ImportRequest, taskSlot int64) error {
ret := _m.Called(nodeID, in, taskSlot)
if len(ret) == 0 {
panic("no return value specified for CreateImport")
}
var r0 error
if rf, ok := ret.Get(0).(func(int64, *datapb.ImportRequest, int64) error); ok {
r0 = rf(nodeID, in, taskSlot)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockCluster_CreateImport_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CreateImport'
type MockCluster_CreateImport_Call struct {
*mock.Call
}
// CreateImport is a helper method to define mock.On call
// - nodeID int64
// - in *datapb.ImportRequest
// - taskSlot int64
func (_e *MockCluster_Expecter) CreateImport(nodeID interface{}, in interface{}, taskSlot interface{}) *MockCluster_CreateImport_Call {
return &MockCluster_CreateImport_Call{Call: _e.mock.On("CreateImport", nodeID, in, taskSlot)}
}
func (_c *MockCluster_CreateImport_Call) Run(run func(nodeID int64, in *datapb.ImportRequest, taskSlot int64)) *MockCluster_CreateImport_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64), args[1].(*datapb.ImportRequest), args[2].(int64))
})
return _c
}
func (_c *MockCluster_CreateImport_Call) Return(_a0 error) *MockCluster_CreateImport_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockCluster_CreateImport_Call) RunAndReturn(run func(int64, *datapb.ImportRequest, int64) error) *MockCluster_CreateImport_Call {
_c.Call.Return(run)
return _c
}
// CreateIndex provides a mock function with given fields: nodeID, in
func (_m *MockCluster) CreateIndex(nodeID int64, in *workerpb.CreateJobRequest) error {
ret := _m.Called(nodeID, in)
if len(ret) == 0 {
panic("no return value specified for CreateIndex")
}
var r0 error
if rf, ok := ret.Get(0).(func(int64, *workerpb.CreateJobRequest) error); ok {
r0 = rf(nodeID, in)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockCluster_CreateIndex_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CreateIndex'
type MockCluster_CreateIndex_Call struct {
*mock.Call
}
// CreateIndex is a helper method to define mock.On call
// - nodeID int64
// - in *workerpb.CreateJobRequest
func (_e *MockCluster_Expecter) CreateIndex(nodeID interface{}, in interface{}) *MockCluster_CreateIndex_Call {
return &MockCluster_CreateIndex_Call{Call: _e.mock.On("CreateIndex", nodeID, in)}
}
func (_c *MockCluster_CreateIndex_Call) Run(run func(nodeID int64, in *workerpb.CreateJobRequest)) *MockCluster_CreateIndex_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64), args[1].(*workerpb.CreateJobRequest))
})
return _c
}
func (_c *MockCluster_CreateIndex_Call) Return(_a0 error) *MockCluster_CreateIndex_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockCluster_CreateIndex_Call) RunAndReturn(run func(int64, *workerpb.CreateJobRequest) error) *MockCluster_CreateIndex_Call {
_c.Call.Return(run)
return _c
}
// CreatePreImport provides a mock function with given fields: nodeID, in, taskSlot
func (_m *MockCluster) CreatePreImport(nodeID int64, in *datapb.PreImportRequest, taskSlot int64) error {
ret := _m.Called(nodeID, in, taskSlot)
if len(ret) == 0 {
panic("no return value specified for CreatePreImport")
}
var r0 error
if rf, ok := ret.Get(0).(func(int64, *datapb.PreImportRequest, int64) error); ok {
r0 = rf(nodeID, in, taskSlot)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockCluster_CreatePreImport_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CreatePreImport'
type MockCluster_CreatePreImport_Call struct {
*mock.Call
}
// CreatePreImport is a helper method to define mock.On call
// - nodeID int64
// - in *datapb.PreImportRequest
// - taskSlot int64
func (_e *MockCluster_Expecter) CreatePreImport(nodeID interface{}, in interface{}, taskSlot interface{}) *MockCluster_CreatePreImport_Call {
return &MockCluster_CreatePreImport_Call{Call: _e.mock.On("CreatePreImport", nodeID, in, taskSlot)}
}
func (_c *MockCluster_CreatePreImport_Call) Run(run func(nodeID int64, in *datapb.PreImportRequest, taskSlot int64)) *MockCluster_CreatePreImport_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64), args[1].(*datapb.PreImportRequest), args[2].(int64))
})
return _c
}
func (_c *MockCluster_CreatePreImport_Call) Return(_a0 error) *MockCluster_CreatePreImport_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockCluster_CreatePreImport_Call) RunAndReturn(run func(int64, *datapb.PreImportRequest, int64) error) *MockCluster_CreatePreImport_Call {
_c.Call.Return(run)
return _c
}
// CreateStats provides a mock function with given fields: nodeID, in
func (_m *MockCluster) CreateStats(nodeID int64, in *workerpb.CreateStatsRequest) error {
ret := _m.Called(nodeID, in)
if len(ret) == 0 {
panic("no return value specified for CreateStats")
}
var r0 error
if rf, ok := ret.Get(0).(func(int64, *workerpb.CreateStatsRequest) error); ok {
r0 = rf(nodeID, in)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockCluster_CreateStats_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CreateStats'
type MockCluster_CreateStats_Call struct {
*mock.Call
}
// CreateStats is a helper method to define mock.On call
// - nodeID int64
// - in *workerpb.CreateStatsRequest
func (_e *MockCluster_Expecter) CreateStats(nodeID interface{}, in interface{}) *MockCluster_CreateStats_Call {
return &MockCluster_CreateStats_Call{Call: _e.mock.On("CreateStats", nodeID, in)}
}
func (_c *MockCluster_CreateStats_Call) Run(run func(nodeID int64, in *workerpb.CreateStatsRequest)) *MockCluster_CreateStats_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64), args[1].(*workerpb.CreateStatsRequest))
})
return _c
}
func (_c *MockCluster_CreateStats_Call) Return(_a0 error) *MockCluster_CreateStats_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockCluster_CreateStats_Call) RunAndReturn(run func(int64, *workerpb.CreateStatsRequest) error) *MockCluster_CreateStats_Call {
_c.Call.Return(run)
return _c
}
// DropAnalyze provides a mock function with given fields: nodeID, taskID
func (_m *MockCluster) DropAnalyze(nodeID int64, taskID int64) error {
ret := _m.Called(nodeID, taskID)
if len(ret) == 0 {
panic("no return value specified for DropAnalyze")
}
var r0 error
if rf, ok := ret.Get(0).(func(int64, int64) error); ok {
r0 = rf(nodeID, taskID)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockCluster_DropAnalyze_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'DropAnalyze'
type MockCluster_DropAnalyze_Call struct {
*mock.Call
}
// DropAnalyze is a helper method to define mock.On call
// - nodeID int64
// - taskID int64
func (_e *MockCluster_Expecter) DropAnalyze(nodeID interface{}, taskID interface{}) *MockCluster_DropAnalyze_Call {
return &MockCluster_DropAnalyze_Call{Call: _e.mock.On("DropAnalyze", nodeID, taskID)}
}
func (_c *MockCluster_DropAnalyze_Call) Run(run func(nodeID int64, taskID int64)) *MockCluster_DropAnalyze_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64), args[1].(int64))
})
return _c
}
func (_c *MockCluster_DropAnalyze_Call) Return(_a0 error) *MockCluster_DropAnalyze_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockCluster_DropAnalyze_Call) RunAndReturn(run func(int64, int64) error) *MockCluster_DropAnalyze_Call {
_c.Call.Return(run)
return _c
}
// DropCompaction provides a mock function with given fields: nodeID, planID
func (_m *MockCluster) DropCompaction(nodeID int64, planID int64) error {
ret := _m.Called(nodeID, planID)
if len(ret) == 0 {
panic("no return value specified for DropCompaction")
}
var r0 error
if rf, ok := ret.Get(0).(func(int64, int64) error); ok {
r0 = rf(nodeID, planID)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockCluster_DropCompaction_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'DropCompaction'
type MockCluster_DropCompaction_Call struct {
*mock.Call
}
// DropCompaction is a helper method to define mock.On call
// - nodeID int64
// - planID int64
func (_e *MockCluster_Expecter) DropCompaction(nodeID interface{}, planID interface{}) *MockCluster_DropCompaction_Call {
return &MockCluster_DropCompaction_Call{Call: _e.mock.On("DropCompaction", nodeID, planID)}
}
func (_c *MockCluster_DropCompaction_Call) Run(run func(nodeID int64, planID int64)) *MockCluster_DropCompaction_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64), args[1].(int64))
})
return _c
}
func (_c *MockCluster_DropCompaction_Call) Return(_a0 error) *MockCluster_DropCompaction_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockCluster_DropCompaction_Call) RunAndReturn(run func(int64, int64) error) *MockCluster_DropCompaction_Call {
_c.Call.Return(run)
return _c
}
// DropExternalCollectionTask provides a mock function with given fields: nodeID, taskID
func (_m *MockCluster) DropExternalCollectionTask(nodeID int64, taskID int64) error {
ret := _m.Called(nodeID, taskID)
if len(ret) == 0 {
panic("no return value specified for DropExternalCollectionTask")
}
var r0 error
if rf, ok := ret.Get(0).(func(int64, int64) error); ok {
r0 = rf(nodeID, taskID)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockCluster_DropExternalCollectionTask_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'DropExternalCollectionTask'
type MockCluster_DropExternalCollectionTask_Call struct {
*mock.Call
}
// DropExternalCollectionTask is a helper method to define mock.On call
// - nodeID int64
// - taskID int64
func (_e *MockCluster_Expecter) DropExternalCollectionTask(nodeID interface{}, taskID interface{}) *MockCluster_DropExternalCollectionTask_Call {
return &MockCluster_DropExternalCollectionTask_Call{Call: _e.mock.On("DropExternalCollectionTask", nodeID, taskID)}
}
func (_c *MockCluster_DropExternalCollectionTask_Call) Run(run func(nodeID int64, taskID int64)) *MockCluster_DropExternalCollectionTask_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64), args[1].(int64))
})
return _c
}
func (_c *MockCluster_DropExternalCollectionTask_Call) Return(_a0 error) *MockCluster_DropExternalCollectionTask_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockCluster_DropExternalCollectionTask_Call) RunAndReturn(run func(int64, int64) error) *MockCluster_DropExternalCollectionTask_Call {
_c.Call.Return(run)
return _c
}
// DropImport provides a mock function with given fields: nodeID, taskID
func (_m *MockCluster) DropImport(nodeID int64, taskID int64) error {
ret := _m.Called(nodeID, taskID)
if len(ret) == 0 {
panic("no return value specified for DropImport")
}
var r0 error
if rf, ok := ret.Get(0).(func(int64, int64) error); ok {
r0 = rf(nodeID, taskID)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockCluster_DropImport_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'DropImport'
type MockCluster_DropImport_Call struct {
*mock.Call
}
// DropImport is a helper method to define mock.On call
// - nodeID int64
// - taskID int64
func (_e *MockCluster_Expecter) DropImport(nodeID interface{}, taskID interface{}) *MockCluster_DropImport_Call {
return &MockCluster_DropImport_Call{Call: _e.mock.On("DropImport", nodeID, taskID)}
}
func (_c *MockCluster_DropImport_Call) Run(run func(nodeID int64, taskID int64)) *MockCluster_DropImport_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64), args[1].(int64))
})
return _c
}
func (_c *MockCluster_DropImport_Call) Return(_a0 error) *MockCluster_DropImport_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockCluster_DropImport_Call) RunAndReturn(run func(int64, int64) error) *MockCluster_DropImport_Call {
_c.Call.Return(run)
return _c
}
// DropIndex provides a mock function with given fields: nodeID, taskID
func (_m *MockCluster) DropIndex(nodeID int64, taskID int64) error {
ret := _m.Called(nodeID, taskID)
if len(ret) == 0 {
panic("no return value specified for DropIndex")
}
var r0 error
if rf, ok := ret.Get(0).(func(int64, int64) error); ok {
r0 = rf(nodeID, taskID)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockCluster_DropIndex_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'DropIndex'
type MockCluster_DropIndex_Call struct {
*mock.Call
}
// DropIndex is a helper method to define mock.On call
// - nodeID int64
// - taskID int64
func (_e *MockCluster_Expecter) DropIndex(nodeID interface{}, taskID interface{}) *MockCluster_DropIndex_Call {
return &MockCluster_DropIndex_Call{Call: _e.mock.On("DropIndex", nodeID, taskID)}
}
func (_c *MockCluster_DropIndex_Call) Run(run func(nodeID int64, taskID int64)) *MockCluster_DropIndex_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64), args[1].(int64))
})
return _c
}
func (_c *MockCluster_DropIndex_Call) Return(_a0 error) *MockCluster_DropIndex_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockCluster_DropIndex_Call) RunAndReturn(run func(int64, int64) error) *MockCluster_DropIndex_Call {
_c.Call.Return(run)
return _c
}
// DropStats provides a mock function with given fields: nodeID, taskID
func (_m *MockCluster) DropStats(nodeID int64, taskID int64) error {
ret := _m.Called(nodeID, taskID)
if len(ret) == 0 {
panic("no return value specified for DropStats")
}
var r0 error
if rf, ok := ret.Get(0).(func(int64, int64) error); ok {
r0 = rf(nodeID, taskID)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockCluster_DropStats_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'DropStats'
type MockCluster_DropStats_Call struct {
*mock.Call
}
// DropStats is a helper method to define mock.On call
// - nodeID int64
// - taskID int64
func (_e *MockCluster_Expecter) DropStats(nodeID interface{}, taskID interface{}) *MockCluster_DropStats_Call {
return &MockCluster_DropStats_Call{Call: _e.mock.On("DropStats", nodeID, taskID)}
}
func (_c *MockCluster_DropStats_Call) Run(run func(nodeID int64, taskID int64)) *MockCluster_DropStats_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64), args[1].(int64))
})
return _c
}
func (_c *MockCluster_DropStats_Call) Return(_a0 error) *MockCluster_DropStats_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockCluster_DropStats_Call) RunAndReturn(run func(int64, int64) error) *MockCluster_DropStats_Call {
_c.Call.Return(run)
return _c
}
// QueryAnalyze provides a mock function with given fields: nodeID, in
func (_m *MockCluster) QueryAnalyze(nodeID int64, in *workerpb.QueryJobsRequest) (*workerpb.AnalyzeResults, error) {
ret := _m.Called(nodeID, in)
if len(ret) == 0 {
panic("no return value specified for QueryAnalyze")
}
var r0 *workerpb.AnalyzeResults
var r1 error
if rf, ok := ret.Get(0).(func(int64, *workerpb.QueryJobsRequest) (*workerpb.AnalyzeResults, error)); ok {
return rf(nodeID, in)
}
if rf, ok := ret.Get(0).(func(int64, *workerpb.QueryJobsRequest) *workerpb.AnalyzeResults); ok {
r0 = rf(nodeID, in)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(*workerpb.AnalyzeResults)
}
}
if rf, ok := ret.Get(1).(func(int64, *workerpb.QueryJobsRequest) error); ok {
r1 = rf(nodeID, in)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MockCluster_QueryAnalyze_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'QueryAnalyze'
type MockCluster_QueryAnalyze_Call struct {
*mock.Call
}
// QueryAnalyze is a helper method to define mock.On call
// - nodeID int64
// - in *workerpb.QueryJobsRequest
func (_e *MockCluster_Expecter) QueryAnalyze(nodeID interface{}, in interface{}) *MockCluster_QueryAnalyze_Call {
return &MockCluster_QueryAnalyze_Call{Call: _e.mock.On("QueryAnalyze", nodeID, in)}
}
func (_c *MockCluster_QueryAnalyze_Call) Run(run func(nodeID int64, in *workerpb.QueryJobsRequest)) *MockCluster_QueryAnalyze_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64), args[1].(*workerpb.QueryJobsRequest))
})
return _c
}
func (_c *MockCluster_QueryAnalyze_Call) Return(_a0 *workerpb.AnalyzeResults, _a1 error) *MockCluster_QueryAnalyze_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockCluster_QueryAnalyze_Call) RunAndReturn(run func(int64, *workerpb.QueryJobsRequest) (*workerpb.AnalyzeResults, error)) *MockCluster_QueryAnalyze_Call {
_c.Call.Return(run)
return _c
}
// QueryCompaction provides a mock function with given fields: nodeID, in
func (_m *MockCluster) QueryCompaction(nodeID int64, in *datapb.CompactionStateRequest) (*datapb.CompactionPlanResult, error) {
ret := _m.Called(nodeID, in)
if len(ret) == 0 {
panic("no return value specified for QueryCompaction")
}
var r0 *datapb.CompactionPlanResult
var r1 error
if rf, ok := ret.Get(0).(func(int64, *datapb.CompactionStateRequest) (*datapb.CompactionPlanResult, error)); ok {
return rf(nodeID, in)
}
if rf, ok := ret.Get(0).(func(int64, *datapb.CompactionStateRequest) *datapb.CompactionPlanResult); ok {
r0 = rf(nodeID, in)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(*datapb.CompactionPlanResult)
}
}
if rf, ok := ret.Get(1).(func(int64, *datapb.CompactionStateRequest) error); ok {
r1 = rf(nodeID, in)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MockCluster_QueryCompaction_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'QueryCompaction'
type MockCluster_QueryCompaction_Call struct {
*mock.Call
}
// QueryCompaction is a helper method to define mock.On call
// - nodeID int64
// - in *datapb.CompactionStateRequest
func (_e *MockCluster_Expecter) QueryCompaction(nodeID interface{}, in interface{}) *MockCluster_QueryCompaction_Call {
return &MockCluster_QueryCompaction_Call{Call: _e.mock.On("QueryCompaction", nodeID, in)}
}
func (_c *MockCluster_QueryCompaction_Call) Run(run func(nodeID int64, in *datapb.CompactionStateRequest)) *MockCluster_QueryCompaction_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64), args[1].(*datapb.CompactionStateRequest))
})
return _c
}
func (_c *MockCluster_QueryCompaction_Call) Return(_a0 *datapb.CompactionPlanResult, _a1 error) *MockCluster_QueryCompaction_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockCluster_QueryCompaction_Call) RunAndReturn(run func(int64, *datapb.CompactionStateRequest) (*datapb.CompactionPlanResult, error)) *MockCluster_QueryCompaction_Call {
_c.Call.Return(run)
return _c
}
// QueryExternalCollectionTask provides a mock function with given fields: nodeID, taskID
func (_m *MockCluster) QueryExternalCollectionTask(nodeID int64, taskID int64) (*datapb.UpdateExternalCollectionResponse, error) {
ret := _m.Called(nodeID, taskID)
if len(ret) == 0 {
panic("no return value specified for QueryExternalCollectionTask")
}
var r0 *datapb.UpdateExternalCollectionResponse
var r1 error
if rf, ok := ret.Get(0).(func(int64, int64) (*datapb.UpdateExternalCollectionResponse, error)); ok {
return rf(nodeID, taskID)
}
if rf, ok := ret.Get(0).(func(int64, int64) *datapb.UpdateExternalCollectionResponse); ok {
r0 = rf(nodeID, taskID)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(*datapb.UpdateExternalCollectionResponse)
}
}
if rf, ok := ret.Get(1).(func(int64, int64) error); ok {
r1 = rf(nodeID, taskID)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MockCluster_QueryExternalCollectionTask_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'QueryExternalCollectionTask'
type MockCluster_QueryExternalCollectionTask_Call struct {
*mock.Call
}
// QueryExternalCollectionTask is a helper method to define mock.On call
// - nodeID int64
// - taskID int64
func (_e *MockCluster_Expecter) QueryExternalCollectionTask(nodeID interface{}, taskID interface{}) *MockCluster_QueryExternalCollectionTask_Call {
return &MockCluster_QueryExternalCollectionTask_Call{Call: _e.mock.On("QueryExternalCollectionTask", nodeID, taskID)}
}
func (_c *MockCluster_QueryExternalCollectionTask_Call) Run(run func(nodeID int64, taskID int64)) *MockCluster_QueryExternalCollectionTask_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64), args[1].(int64))
})
return _c
}
func (_c *MockCluster_QueryExternalCollectionTask_Call) Return(_a0 *datapb.UpdateExternalCollectionResponse, _a1 error) *MockCluster_QueryExternalCollectionTask_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockCluster_QueryExternalCollectionTask_Call) RunAndReturn(run func(int64, int64) (*datapb.UpdateExternalCollectionResponse, error)) *MockCluster_QueryExternalCollectionTask_Call {
_c.Call.Return(run)
return _c
}
// QueryImport provides a mock function with given fields: nodeID, in
func (_m *MockCluster) QueryImport(nodeID int64, in *datapb.QueryImportRequest) (*datapb.QueryImportResponse, error) {
ret := _m.Called(nodeID, in)
if len(ret) == 0 {
panic("no return value specified for QueryImport")
}
var r0 *datapb.QueryImportResponse
var r1 error
if rf, ok := ret.Get(0).(func(int64, *datapb.QueryImportRequest) (*datapb.QueryImportResponse, error)); ok {
return rf(nodeID, in)
}
if rf, ok := ret.Get(0).(func(int64, *datapb.QueryImportRequest) *datapb.QueryImportResponse); ok {
r0 = rf(nodeID, in)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(*datapb.QueryImportResponse)
}
}
if rf, ok := ret.Get(1).(func(int64, *datapb.QueryImportRequest) error); ok {
r1 = rf(nodeID, in)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MockCluster_QueryImport_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'QueryImport'
type MockCluster_QueryImport_Call struct {
*mock.Call
}
// QueryImport is a helper method to define mock.On call
// - nodeID int64
// - in *datapb.QueryImportRequest
func (_e *MockCluster_Expecter) QueryImport(nodeID interface{}, in interface{}) *MockCluster_QueryImport_Call {
return &MockCluster_QueryImport_Call{Call: _e.mock.On("QueryImport", nodeID, in)}
}
func (_c *MockCluster_QueryImport_Call) Run(run func(nodeID int64, in *datapb.QueryImportRequest)) *MockCluster_QueryImport_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64), args[1].(*datapb.QueryImportRequest))
})
return _c
}
func (_c *MockCluster_QueryImport_Call) Return(_a0 *datapb.QueryImportResponse, _a1 error) *MockCluster_QueryImport_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockCluster_QueryImport_Call) RunAndReturn(run func(int64, *datapb.QueryImportRequest) (*datapb.QueryImportResponse, error)) *MockCluster_QueryImport_Call {
_c.Call.Return(run)
return _c
}
// QueryIndex provides a mock function with given fields: nodeID, in
func (_m *MockCluster) QueryIndex(nodeID int64, in *workerpb.QueryJobsRequest) (*workerpb.IndexJobResults, error) {
ret := _m.Called(nodeID, in)
if len(ret) == 0 {
panic("no return value specified for QueryIndex")
}
var r0 *workerpb.IndexJobResults
var r1 error
if rf, ok := ret.Get(0).(func(int64, *workerpb.QueryJobsRequest) (*workerpb.IndexJobResults, error)); ok {
return rf(nodeID, in)
}
if rf, ok := ret.Get(0).(func(int64, *workerpb.QueryJobsRequest) *workerpb.IndexJobResults); ok {
r0 = rf(nodeID, in)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(*workerpb.IndexJobResults)
}
}
if rf, ok := ret.Get(1).(func(int64, *workerpb.QueryJobsRequest) error); ok {
r1 = rf(nodeID, in)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MockCluster_QueryIndex_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'QueryIndex'
type MockCluster_QueryIndex_Call struct {
*mock.Call
}
// QueryIndex is a helper method to define mock.On call
// - nodeID int64
// - in *workerpb.QueryJobsRequest
func (_e *MockCluster_Expecter) QueryIndex(nodeID interface{}, in interface{}) *MockCluster_QueryIndex_Call {
return &MockCluster_QueryIndex_Call{Call: _e.mock.On("QueryIndex", nodeID, in)}
}
func (_c *MockCluster_QueryIndex_Call) Run(run func(nodeID int64, in *workerpb.QueryJobsRequest)) *MockCluster_QueryIndex_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64), args[1].(*workerpb.QueryJobsRequest))
})
return _c
}
func (_c *MockCluster_QueryIndex_Call) Return(_a0 *workerpb.IndexJobResults, _a1 error) *MockCluster_QueryIndex_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockCluster_QueryIndex_Call) RunAndReturn(run func(int64, *workerpb.QueryJobsRequest) (*workerpb.IndexJobResults, error)) *MockCluster_QueryIndex_Call {
_c.Call.Return(run)
return _c
}
// QueryPreImport provides a mock function with given fields: nodeID, in
func (_m *MockCluster) QueryPreImport(nodeID int64, in *datapb.QueryPreImportRequest) (*datapb.QueryPreImportResponse, error) {
ret := _m.Called(nodeID, in)
if len(ret) == 0 {
panic("no return value specified for QueryPreImport")
}
var r0 *datapb.QueryPreImportResponse
var r1 error
if rf, ok := ret.Get(0).(func(int64, *datapb.QueryPreImportRequest) (*datapb.QueryPreImportResponse, error)); ok {
return rf(nodeID, in)
}
if rf, ok := ret.Get(0).(func(int64, *datapb.QueryPreImportRequest) *datapb.QueryPreImportResponse); ok {
r0 = rf(nodeID, in)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(*datapb.QueryPreImportResponse)
}
}
if rf, ok := ret.Get(1).(func(int64, *datapb.QueryPreImportRequest) error); ok {
r1 = rf(nodeID, in)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MockCluster_QueryPreImport_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'QueryPreImport'
type MockCluster_QueryPreImport_Call struct {
*mock.Call
}
// QueryPreImport is a helper method to define mock.On call
// - nodeID int64
// - in *datapb.QueryPreImportRequest
func (_e *MockCluster_Expecter) QueryPreImport(nodeID interface{}, in interface{}) *MockCluster_QueryPreImport_Call {
return &MockCluster_QueryPreImport_Call{Call: _e.mock.On("QueryPreImport", nodeID, in)}
}
func (_c *MockCluster_QueryPreImport_Call) Run(run func(nodeID int64, in *datapb.QueryPreImportRequest)) *MockCluster_QueryPreImport_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64), args[1].(*datapb.QueryPreImportRequest))
})
return _c
}
func (_c *MockCluster_QueryPreImport_Call) Return(_a0 *datapb.QueryPreImportResponse, _a1 error) *MockCluster_QueryPreImport_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockCluster_QueryPreImport_Call) RunAndReturn(run func(int64, *datapb.QueryPreImportRequest) (*datapb.QueryPreImportResponse, error)) *MockCluster_QueryPreImport_Call {
_c.Call.Return(run)
return _c
}
// QuerySlot provides a mock function with no fields
func (_m *MockCluster) QuerySlot() map[int64]*WorkerSlots {
ret := _m.Called()
if len(ret) == 0 {
panic("no return value specified for QuerySlot")
}
var r0 map[int64]*WorkerSlots
if rf, ok := ret.Get(0).(func() map[int64]*WorkerSlots); ok {
r0 = rf()
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(map[int64]*WorkerSlots)
}
}
return r0
}
// MockCluster_QuerySlot_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'QuerySlot'
type MockCluster_QuerySlot_Call struct {
*mock.Call
}
// QuerySlot is a helper method to define mock.On call
func (_e *MockCluster_Expecter) QuerySlot() *MockCluster_QuerySlot_Call {
return &MockCluster_QuerySlot_Call{Call: _e.mock.On("QuerySlot")}
}
func (_c *MockCluster_QuerySlot_Call) Run(run func()) *MockCluster_QuerySlot_Call {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *MockCluster_QuerySlot_Call) Return(_a0 map[int64]*WorkerSlots) *MockCluster_QuerySlot_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockCluster_QuerySlot_Call) RunAndReturn(run func() map[int64]*WorkerSlots) *MockCluster_QuerySlot_Call {
_c.Call.Return(run)
return _c
}
// QueryStats provides a mock function with given fields: nodeID, in
func (_m *MockCluster) QueryStats(nodeID int64, in *workerpb.QueryJobsRequest) (*workerpb.StatsResults, error) {
ret := _m.Called(nodeID, in)
if len(ret) == 0 {
panic("no return value specified for QueryStats")
}
var r0 *workerpb.StatsResults
var r1 error
if rf, ok := ret.Get(0).(func(int64, *workerpb.QueryJobsRequest) (*workerpb.StatsResults, error)); ok {
return rf(nodeID, in)
}
if rf, ok := ret.Get(0).(func(int64, *workerpb.QueryJobsRequest) *workerpb.StatsResults); ok {
r0 = rf(nodeID, in)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(*workerpb.StatsResults)
}
}
if rf, ok := ret.Get(1).(func(int64, *workerpb.QueryJobsRequest) error); ok {
r1 = rf(nodeID, in)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MockCluster_QueryStats_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'QueryStats'
type MockCluster_QueryStats_Call struct {
*mock.Call
}
// QueryStats is a helper method to define mock.On call
// - nodeID int64
// - in *workerpb.QueryJobsRequest
func (_e *MockCluster_Expecter) QueryStats(nodeID interface{}, in interface{}) *MockCluster_QueryStats_Call {
return &MockCluster_QueryStats_Call{Call: _e.mock.On("QueryStats", nodeID, in)}
}
func (_c *MockCluster_QueryStats_Call) Run(run func(nodeID int64, in *workerpb.QueryJobsRequest)) *MockCluster_QueryStats_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64), args[1].(*workerpb.QueryJobsRequest))
})
return _c
}
func (_c *MockCluster_QueryStats_Call) Return(_a0 *workerpb.StatsResults, _a1 error) *MockCluster_QueryStats_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockCluster_QueryStats_Call) RunAndReturn(run func(int64, *workerpb.QueryJobsRequest) (*workerpb.StatsResults, error)) *MockCluster_QueryStats_Call {
_c.Call.Return(run)
return _c
}
// NewMockCluster creates a new instance of MockCluster. 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 NewMockCluster(t interface {
mock.TestingT
Cleanup(func())
}) *MockCluster {
mock := &MockCluster{}
mock.Mock.Test(t)
t.Cleanup(func() { mock.AssertExpectations(t) })
return mock
}