milvus/internal/datacoord/session/mock_datanode_manager.go
jaime 28fdbc4e30
enhance: optimize CPU usage for CheckHealth requests (#35589)
issue: #35563
1. Use an internal health checker to monitor the cluster's health state,
storing the latest state on the coordinator node. The CheckHealth
request retrieves the cluster's health from this latest state on the
proxy sides, which enhances cluster stability.
2. Each health check will assess all collections and channels, with
detailed failure messages temporarily saved in the latest state.
3. Use CheckHealth request instead of the heavy GetMetrics request on
the querynode and datanode

Signed-off-by: jaime <yun.zhang@zilliz.com>
2024-12-17 11:02:45 +08:00

1107 lines
37 KiB
Go

// Code generated by mockery v2.46.0. DO NOT EDIT.
package session
import (
context "context"
datapb "github.com/milvus-io/milvus/internal/proto/datapb"
healthcheck "github.com/milvus-io/milvus/internal/util/healthcheck"
mock "github.com/stretchr/testify/mock"
typeutil "github.com/milvus-io/milvus/pkg/util/typeutil"
)
// MockDataNodeManager is an autogenerated mock type for the DataNodeManager type
type MockDataNodeManager struct {
mock.Mock
}
type MockDataNodeManager_Expecter struct {
mock *mock.Mock
}
func (_m *MockDataNodeManager) EXPECT() *MockDataNodeManager_Expecter {
return &MockDataNodeManager_Expecter{mock: &_m.Mock}
}
// AddSession provides a mock function with given fields: node
func (_m *MockDataNodeManager) AddSession(node *NodeInfo) {
_m.Called(node)
}
// MockDataNodeManager_AddSession_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'AddSession'
type MockDataNodeManager_AddSession_Call struct {
*mock.Call
}
// AddSession is a helper method to define mock.On call
// - node *NodeInfo
func (_e *MockDataNodeManager_Expecter) AddSession(node interface{}) *MockDataNodeManager_AddSession_Call {
return &MockDataNodeManager_AddSession_Call{Call: _e.mock.On("AddSession", node)}
}
func (_c *MockDataNodeManager_AddSession_Call) Run(run func(node *NodeInfo)) *MockDataNodeManager_AddSession_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(*NodeInfo))
})
return _c
}
func (_c *MockDataNodeManager_AddSession_Call) Return() *MockDataNodeManager_AddSession_Call {
_c.Call.Return()
return _c
}
func (_c *MockDataNodeManager_AddSession_Call) RunAndReturn(run func(*NodeInfo)) *MockDataNodeManager_AddSession_Call {
_c.Call.Return(run)
return _c
}
// CheckChannelOperationProgress provides a mock function with given fields: ctx, nodeID, info
func (_m *MockDataNodeManager) CheckChannelOperationProgress(ctx context.Context, nodeID int64, info *datapb.ChannelWatchInfo) (*datapb.ChannelOperationProgressResponse, error) {
ret := _m.Called(ctx, nodeID, info)
if len(ret) == 0 {
panic("no return value specified for CheckChannelOperationProgress")
}
var r0 *datapb.ChannelOperationProgressResponse
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, int64, *datapb.ChannelWatchInfo) (*datapb.ChannelOperationProgressResponse, error)); ok {
return rf(ctx, nodeID, info)
}
if rf, ok := ret.Get(0).(func(context.Context, int64, *datapb.ChannelWatchInfo) *datapb.ChannelOperationProgressResponse); ok {
r0 = rf(ctx, nodeID, info)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(*datapb.ChannelOperationProgressResponse)
}
}
if rf, ok := ret.Get(1).(func(context.Context, int64, *datapb.ChannelWatchInfo) error); ok {
r1 = rf(ctx, nodeID, info)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MockDataNodeManager_CheckChannelOperationProgress_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CheckChannelOperationProgress'
type MockDataNodeManager_CheckChannelOperationProgress_Call struct {
*mock.Call
}
// CheckChannelOperationProgress is a helper method to define mock.On call
// - ctx context.Context
// - nodeID int64
// - info *datapb.ChannelWatchInfo
func (_e *MockDataNodeManager_Expecter) CheckChannelOperationProgress(ctx interface{}, nodeID interface{}, info interface{}) *MockDataNodeManager_CheckChannelOperationProgress_Call {
return &MockDataNodeManager_CheckChannelOperationProgress_Call{Call: _e.mock.On("CheckChannelOperationProgress", ctx, nodeID, info)}
}
func (_c *MockDataNodeManager_CheckChannelOperationProgress_Call) Run(run func(ctx context.Context, nodeID int64, info *datapb.ChannelWatchInfo)) *MockDataNodeManager_CheckChannelOperationProgress_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(int64), args[2].(*datapb.ChannelWatchInfo))
})
return _c
}
func (_c *MockDataNodeManager_CheckChannelOperationProgress_Call) Return(_a0 *datapb.ChannelOperationProgressResponse, _a1 error) *MockDataNodeManager_CheckChannelOperationProgress_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockDataNodeManager_CheckChannelOperationProgress_Call) RunAndReturn(run func(context.Context, int64, *datapb.ChannelWatchInfo) (*datapb.ChannelOperationProgressResponse, error)) *MockDataNodeManager_CheckChannelOperationProgress_Call {
_c.Call.Return(run)
return _c
}
// CheckDNHealth provides a mock function with given fields: ctx
func (_m *MockDataNodeManager) CheckDNHealth(ctx context.Context) *healthcheck.Result {
ret := _m.Called(ctx)
if len(ret) == 0 {
panic("no return value specified for CheckDNHealth")
}
var r0 *healthcheck.Result
if rf, ok := ret.Get(0).(func(context.Context) *healthcheck.Result); ok {
r0 = rf(ctx)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(*healthcheck.Result)
}
}
return r0
}
// MockDataNodeManager_CheckDNHealth_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CheckDNHealth'
type MockDataNodeManager_CheckDNHealth_Call struct {
*mock.Call
}
// CheckDNHealth is a helper method to define mock.On call
// - ctx context.Context
func (_e *MockDataNodeManager_Expecter) CheckDNHealth(ctx interface{}) *MockDataNodeManager_CheckDNHealth_Call {
return &MockDataNodeManager_CheckDNHealth_Call{Call: _e.mock.On("CheckDNHealth", ctx)}
}
func (_c *MockDataNodeManager_CheckDNHealth_Call) Run(run func(ctx context.Context)) *MockDataNodeManager_CheckDNHealth_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context))
})
return _c
}
func (_c *MockDataNodeManager_CheckDNHealth_Call) Return(_a0 *healthcheck.Result) *MockDataNodeManager_CheckDNHealth_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockDataNodeManager_CheckDNHealth_Call) RunAndReturn(run func(context.Context) *healthcheck.Result) *MockDataNodeManager_CheckDNHealth_Call {
_c.Call.Return(run)
return _c
}
// Close provides a mock function with given fields:
func (_m *MockDataNodeManager) Close() {
_m.Called()
}
// MockDataNodeManager_Close_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Close'
type MockDataNodeManager_Close_Call struct {
*mock.Call
}
// Close is a helper method to define mock.On call
func (_e *MockDataNodeManager_Expecter) Close() *MockDataNodeManager_Close_Call {
return &MockDataNodeManager_Close_Call{Call: _e.mock.On("Close")}
}
func (_c *MockDataNodeManager_Close_Call) Run(run func()) *MockDataNodeManager_Close_Call {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *MockDataNodeManager_Close_Call) Return() *MockDataNodeManager_Close_Call {
_c.Call.Return()
return _c
}
func (_c *MockDataNodeManager_Close_Call) RunAndReturn(run func()) *MockDataNodeManager_Close_Call {
_c.Call.Return(run)
return _c
}
// Compaction provides a mock function with given fields: ctx, nodeID, plan
func (_m *MockDataNodeManager) Compaction(ctx context.Context, nodeID int64, plan *datapb.CompactionPlan) error {
ret := _m.Called(ctx, nodeID, plan)
if len(ret) == 0 {
panic("no return value specified for Compaction")
}
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, int64, *datapb.CompactionPlan) error); ok {
r0 = rf(ctx, nodeID, plan)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockDataNodeManager_Compaction_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Compaction'
type MockDataNodeManager_Compaction_Call struct {
*mock.Call
}
// Compaction is a helper method to define mock.On call
// - ctx context.Context
// - nodeID int64
// - plan *datapb.CompactionPlan
func (_e *MockDataNodeManager_Expecter) Compaction(ctx interface{}, nodeID interface{}, plan interface{}) *MockDataNodeManager_Compaction_Call {
return &MockDataNodeManager_Compaction_Call{Call: _e.mock.On("Compaction", ctx, nodeID, plan)}
}
func (_c *MockDataNodeManager_Compaction_Call) Run(run func(ctx context.Context, nodeID int64, plan *datapb.CompactionPlan)) *MockDataNodeManager_Compaction_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(int64), args[2].(*datapb.CompactionPlan))
})
return _c
}
func (_c *MockDataNodeManager_Compaction_Call) Return(_a0 error) *MockDataNodeManager_Compaction_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockDataNodeManager_Compaction_Call) RunAndReturn(run func(context.Context, int64, *datapb.CompactionPlan) error) *MockDataNodeManager_Compaction_Call {
_c.Call.Return(run)
return _c
}
// DeleteSession provides a mock function with given fields: node
func (_m *MockDataNodeManager) DeleteSession(node *NodeInfo) {
_m.Called(node)
}
// MockDataNodeManager_DeleteSession_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'DeleteSession'
type MockDataNodeManager_DeleteSession_Call struct {
*mock.Call
}
// DeleteSession is a helper method to define mock.On call
// - node *NodeInfo
func (_e *MockDataNodeManager_Expecter) DeleteSession(node interface{}) *MockDataNodeManager_DeleteSession_Call {
return &MockDataNodeManager_DeleteSession_Call{Call: _e.mock.On("DeleteSession", node)}
}
func (_c *MockDataNodeManager_DeleteSession_Call) Run(run func(node *NodeInfo)) *MockDataNodeManager_DeleteSession_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(*NodeInfo))
})
return _c
}
func (_c *MockDataNodeManager_DeleteSession_Call) Return() *MockDataNodeManager_DeleteSession_Call {
_c.Call.Return()
return _c
}
func (_c *MockDataNodeManager_DeleteSession_Call) RunAndReturn(run func(*NodeInfo)) *MockDataNodeManager_DeleteSession_Call {
_c.Call.Return(run)
return _c
}
// DropCompactionPlan provides a mock function with given fields: nodeID, req
func (_m *MockDataNodeManager) DropCompactionPlan(nodeID int64, req *datapb.DropCompactionPlanRequest) error {
ret := _m.Called(nodeID, req)
if len(ret) == 0 {
panic("no return value specified for DropCompactionPlan")
}
var r0 error
if rf, ok := ret.Get(0).(func(int64, *datapb.DropCompactionPlanRequest) error); ok {
r0 = rf(nodeID, req)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockDataNodeManager_DropCompactionPlan_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'DropCompactionPlan'
type MockDataNodeManager_DropCompactionPlan_Call struct {
*mock.Call
}
// DropCompactionPlan is a helper method to define mock.On call
// - nodeID int64
// - req *datapb.DropCompactionPlanRequest
func (_e *MockDataNodeManager_Expecter) DropCompactionPlan(nodeID interface{}, req interface{}) *MockDataNodeManager_DropCompactionPlan_Call {
return &MockDataNodeManager_DropCompactionPlan_Call{Call: _e.mock.On("DropCompactionPlan", nodeID, req)}
}
func (_c *MockDataNodeManager_DropCompactionPlan_Call) Run(run func(nodeID int64, req *datapb.DropCompactionPlanRequest)) *MockDataNodeManager_DropCompactionPlan_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64), args[1].(*datapb.DropCompactionPlanRequest))
})
return _c
}
func (_c *MockDataNodeManager_DropCompactionPlan_Call) Return(_a0 error) *MockDataNodeManager_DropCompactionPlan_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockDataNodeManager_DropCompactionPlan_Call) RunAndReturn(run func(int64, *datapb.DropCompactionPlanRequest) error) *MockDataNodeManager_DropCompactionPlan_Call {
_c.Call.Return(run)
return _c
}
// DropImport provides a mock function with given fields: nodeID, in
func (_m *MockDataNodeManager) DropImport(nodeID int64, in *datapb.DropImportRequest) error {
ret := _m.Called(nodeID, in)
if len(ret) == 0 {
panic("no return value specified for DropImport")
}
var r0 error
if rf, ok := ret.Get(0).(func(int64, *datapb.DropImportRequest) error); ok {
r0 = rf(nodeID, in)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockDataNodeManager_DropImport_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'DropImport'
type MockDataNodeManager_DropImport_Call struct {
*mock.Call
}
// DropImport is a helper method to define mock.On call
// - nodeID int64
// - in *datapb.DropImportRequest
func (_e *MockDataNodeManager_Expecter) DropImport(nodeID interface{}, in interface{}) *MockDataNodeManager_DropImport_Call {
return &MockDataNodeManager_DropImport_Call{Call: _e.mock.On("DropImport", nodeID, in)}
}
func (_c *MockDataNodeManager_DropImport_Call) Run(run func(nodeID int64, in *datapb.DropImportRequest)) *MockDataNodeManager_DropImport_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64), args[1].(*datapb.DropImportRequest))
})
return _c
}
func (_c *MockDataNodeManager_DropImport_Call) Return(_a0 error) *MockDataNodeManager_DropImport_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockDataNodeManager_DropImport_Call) RunAndReturn(run func(int64, *datapb.DropImportRequest) error) *MockDataNodeManager_DropImport_Call {
_c.Call.Return(run)
return _c
}
// Flush provides a mock function with given fields: ctx, nodeID, req
func (_m *MockDataNodeManager) Flush(ctx context.Context, nodeID int64, req *datapb.FlushSegmentsRequest) {
_m.Called(ctx, nodeID, req)
}
// MockDataNodeManager_Flush_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Flush'
type MockDataNodeManager_Flush_Call struct {
*mock.Call
}
// Flush is a helper method to define mock.On call
// - ctx context.Context
// - nodeID int64
// - req *datapb.FlushSegmentsRequest
func (_e *MockDataNodeManager_Expecter) Flush(ctx interface{}, nodeID interface{}, req interface{}) *MockDataNodeManager_Flush_Call {
return &MockDataNodeManager_Flush_Call{Call: _e.mock.On("Flush", ctx, nodeID, req)}
}
func (_c *MockDataNodeManager_Flush_Call) Run(run func(ctx context.Context, nodeID int64, req *datapb.FlushSegmentsRequest)) *MockDataNodeManager_Flush_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(int64), args[2].(*datapb.FlushSegmentsRequest))
})
return _c
}
func (_c *MockDataNodeManager_Flush_Call) Return() *MockDataNodeManager_Flush_Call {
_c.Call.Return()
return _c
}
func (_c *MockDataNodeManager_Flush_Call) RunAndReturn(run func(context.Context, int64, *datapb.FlushSegmentsRequest)) *MockDataNodeManager_Flush_Call {
_c.Call.Return(run)
return _c
}
// FlushChannels provides a mock function with given fields: ctx, nodeID, req
func (_m *MockDataNodeManager) FlushChannels(ctx context.Context, nodeID int64, req *datapb.FlushChannelsRequest) error {
ret := _m.Called(ctx, nodeID, req)
if len(ret) == 0 {
panic("no return value specified for FlushChannels")
}
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, int64, *datapb.FlushChannelsRequest) error); ok {
r0 = rf(ctx, nodeID, req)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockDataNodeManager_FlushChannels_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'FlushChannels'
type MockDataNodeManager_FlushChannels_Call struct {
*mock.Call
}
// FlushChannels is a helper method to define mock.On call
// - ctx context.Context
// - nodeID int64
// - req *datapb.FlushChannelsRequest
func (_e *MockDataNodeManager_Expecter) FlushChannels(ctx interface{}, nodeID interface{}, req interface{}) *MockDataNodeManager_FlushChannels_Call {
return &MockDataNodeManager_FlushChannels_Call{Call: _e.mock.On("FlushChannels", ctx, nodeID, req)}
}
func (_c *MockDataNodeManager_FlushChannels_Call) Run(run func(ctx context.Context, nodeID int64, req *datapb.FlushChannelsRequest)) *MockDataNodeManager_FlushChannels_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(int64), args[2].(*datapb.FlushChannelsRequest))
})
return _c
}
func (_c *MockDataNodeManager_FlushChannels_Call) Return(_a0 error) *MockDataNodeManager_FlushChannels_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockDataNodeManager_FlushChannels_Call) RunAndReturn(run func(context.Context, int64, *datapb.FlushChannelsRequest) error) *MockDataNodeManager_FlushChannels_Call {
_c.Call.Return(run)
return _c
}
// GetCompactionPlanResult provides a mock function with given fields: nodeID, planID
func (_m *MockDataNodeManager) GetCompactionPlanResult(nodeID int64, planID int64) (*datapb.CompactionPlanResult, error) {
ret := _m.Called(nodeID, planID)
if len(ret) == 0 {
panic("no return value specified for GetCompactionPlanResult")
}
var r0 *datapb.CompactionPlanResult
var r1 error
if rf, ok := ret.Get(0).(func(int64, int64) (*datapb.CompactionPlanResult, error)); ok {
return rf(nodeID, planID)
}
if rf, ok := ret.Get(0).(func(int64, int64) *datapb.CompactionPlanResult); ok {
r0 = rf(nodeID, planID)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(*datapb.CompactionPlanResult)
}
}
if rf, ok := ret.Get(1).(func(int64, int64) error); ok {
r1 = rf(nodeID, planID)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MockDataNodeManager_GetCompactionPlanResult_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetCompactionPlanResult'
type MockDataNodeManager_GetCompactionPlanResult_Call struct {
*mock.Call
}
// GetCompactionPlanResult is a helper method to define mock.On call
// - nodeID int64
// - planID int64
func (_e *MockDataNodeManager_Expecter) GetCompactionPlanResult(nodeID interface{}, planID interface{}) *MockDataNodeManager_GetCompactionPlanResult_Call {
return &MockDataNodeManager_GetCompactionPlanResult_Call{Call: _e.mock.On("GetCompactionPlanResult", nodeID, planID)}
}
func (_c *MockDataNodeManager_GetCompactionPlanResult_Call) Run(run func(nodeID int64, planID int64)) *MockDataNodeManager_GetCompactionPlanResult_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64), args[1].(int64))
})
return _c
}
func (_c *MockDataNodeManager_GetCompactionPlanResult_Call) Return(_a0 *datapb.CompactionPlanResult, _a1 error) *MockDataNodeManager_GetCompactionPlanResult_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockDataNodeManager_GetCompactionPlanResult_Call) RunAndReturn(run func(int64, int64) (*datapb.CompactionPlanResult, error)) *MockDataNodeManager_GetCompactionPlanResult_Call {
_c.Call.Return(run)
return _c
}
// GetCompactionPlansResults provides a mock function with given fields:
func (_m *MockDataNodeManager) GetCompactionPlansResults() (map[int64]*typeutil.Pair[int64, *datapb.CompactionPlanResult], error) {
ret := _m.Called()
if len(ret) == 0 {
panic("no return value specified for GetCompactionPlansResults")
}
var r0 map[int64]*typeutil.Pair[int64, *datapb.CompactionPlanResult]
var r1 error
if rf, ok := ret.Get(0).(func() (map[int64]*typeutil.Pair[int64, *datapb.CompactionPlanResult], error)); ok {
return rf()
}
if rf, ok := ret.Get(0).(func() map[int64]*typeutil.Pair[int64, *datapb.CompactionPlanResult]); ok {
r0 = rf()
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(map[int64]*typeutil.Pair[int64, *datapb.CompactionPlanResult])
}
}
if rf, ok := ret.Get(1).(func() error); ok {
r1 = rf()
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MockDataNodeManager_GetCompactionPlansResults_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetCompactionPlansResults'
type MockDataNodeManager_GetCompactionPlansResults_Call struct {
*mock.Call
}
// GetCompactionPlansResults is a helper method to define mock.On call
func (_e *MockDataNodeManager_Expecter) GetCompactionPlansResults() *MockDataNodeManager_GetCompactionPlansResults_Call {
return &MockDataNodeManager_GetCompactionPlansResults_Call{Call: _e.mock.On("GetCompactionPlansResults")}
}
func (_c *MockDataNodeManager_GetCompactionPlansResults_Call) Run(run func()) *MockDataNodeManager_GetCompactionPlansResults_Call {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *MockDataNodeManager_GetCompactionPlansResults_Call) Return(_a0 map[int64]*typeutil.Pair[int64, *datapb.CompactionPlanResult], _a1 error) *MockDataNodeManager_GetCompactionPlansResults_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockDataNodeManager_GetCompactionPlansResults_Call) RunAndReturn(run func() (map[int64]*typeutil.Pair[int64, *datapb.CompactionPlanResult], error)) *MockDataNodeManager_GetCompactionPlansResults_Call {
_c.Call.Return(run)
return _c
}
// GetSession provides a mock function with given fields: _a0
func (_m *MockDataNodeManager) GetSession(_a0 int64) (*Session, bool) {
ret := _m.Called(_a0)
if len(ret) == 0 {
panic("no return value specified for GetSession")
}
var r0 *Session
var r1 bool
if rf, ok := ret.Get(0).(func(int64) (*Session, bool)); ok {
return rf(_a0)
}
if rf, ok := ret.Get(0).(func(int64) *Session); ok {
r0 = rf(_a0)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(*Session)
}
}
if rf, ok := ret.Get(1).(func(int64) bool); ok {
r1 = rf(_a0)
} else {
r1 = ret.Get(1).(bool)
}
return r0, r1
}
// MockDataNodeManager_GetSession_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetSession'
type MockDataNodeManager_GetSession_Call struct {
*mock.Call
}
// GetSession is a helper method to define mock.On call
// - _a0 int64
func (_e *MockDataNodeManager_Expecter) GetSession(_a0 interface{}) *MockDataNodeManager_GetSession_Call {
return &MockDataNodeManager_GetSession_Call{Call: _e.mock.On("GetSession", _a0)}
}
func (_c *MockDataNodeManager_GetSession_Call) Run(run func(_a0 int64)) *MockDataNodeManager_GetSession_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64))
})
return _c
}
func (_c *MockDataNodeManager_GetSession_Call) Return(_a0 *Session, _a1 bool) *MockDataNodeManager_GetSession_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockDataNodeManager_GetSession_Call) RunAndReturn(run func(int64) (*Session, bool)) *MockDataNodeManager_GetSession_Call {
_c.Call.Return(run)
return _c
}
// GetSessionIDs provides a mock function with given fields:
func (_m *MockDataNodeManager) GetSessionIDs() []int64 {
ret := _m.Called()
if len(ret) == 0 {
panic("no return value specified for GetSessionIDs")
}
var r0 []int64
if rf, ok := ret.Get(0).(func() []int64); ok {
r0 = rf()
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).([]int64)
}
}
return r0
}
// MockDataNodeManager_GetSessionIDs_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetSessionIDs'
type MockDataNodeManager_GetSessionIDs_Call struct {
*mock.Call
}
// GetSessionIDs is a helper method to define mock.On call
func (_e *MockDataNodeManager_Expecter) GetSessionIDs() *MockDataNodeManager_GetSessionIDs_Call {
return &MockDataNodeManager_GetSessionIDs_Call{Call: _e.mock.On("GetSessionIDs")}
}
func (_c *MockDataNodeManager_GetSessionIDs_Call) Run(run func()) *MockDataNodeManager_GetSessionIDs_Call {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *MockDataNodeManager_GetSessionIDs_Call) Return(_a0 []int64) *MockDataNodeManager_GetSessionIDs_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockDataNodeManager_GetSessionIDs_Call) RunAndReturn(run func() []int64) *MockDataNodeManager_GetSessionIDs_Call {
_c.Call.Return(run)
return _c
}
// GetSessions provides a mock function with given fields:
func (_m *MockDataNodeManager) GetSessions() []*Session {
ret := _m.Called()
if len(ret) == 0 {
panic("no return value specified for GetSessions")
}
var r0 []*Session
if rf, ok := ret.Get(0).(func() []*Session); ok {
r0 = rf()
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).([]*Session)
}
}
return r0
}
// MockDataNodeManager_GetSessions_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetSessions'
type MockDataNodeManager_GetSessions_Call struct {
*mock.Call
}
// GetSessions is a helper method to define mock.On call
func (_e *MockDataNodeManager_Expecter) GetSessions() *MockDataNodeManager_GetSessions_Call {
return &MockDataNodeManager_GetSessions_Call{Call: _e.mock.On("GetSessions")}
}
func (_c *MockDataNodeManager_GetSessions_Call) Run(run func()) *MockDataNodeManager_GetSessions_Call {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *MockDataNodeManager_GetSessions_Call) Return(_a0 []*Session) *MockDataNodeManager_GetSessions_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockDataNodeManager_GetSessions_Call) RunAndReturn(run func() []*Session) *MockDataNodeManager_GetSessions_Call {
_c.Call.Return(run)
return _c
}
// ImportV2 provides a mock function with given fields: nodeID, in
func (_m *MockDataNodeManager) ImportV2(nodeID int64, in *datapb.ImportRequest) error {
ret := _m.Called(nodeID, in)
if len(ret) == 0 {
panic("no return value specified for ImportV2")
}
var r0 error
if rf, ok := ret.Get(0).(func(int64, *datapb.ImportRequest) error); ok {
r0 = rf(nodeID, in)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockDataNodeManager_ImportV2_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'ImportV2'
type MockDataNodeManager_ImportV2_Call struct {
*mock.Call
}
// ImportV2 is a helper method to define mock.On call
// - nodeID int64
// - in *datapb.ImportRequest
func (_e *MockDataNodeManager_Expecter) ImportV2(nodeID interface{}, in interface{}) *MockDataNodeManager_ImportV2_Call {
return &MockDataNodeManager_ImportV2_Call{Call: _e.mock.On("ImportV2", nodeID, in)}
}
func (_c *MockDataNodeManager_ImportV2_Call) Run(run func(nodeID int64, in *datapb.ImportRequest)) *MockDataNodeManager_ImportV2_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64), args[1].(*datapb.ImportRequest))
})
return _c
}
func (_c *MockDataNodeManager_ImportV2_Call) Return(_a0 error) *MockDataNodeManager_ImportV2_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockDataNodeManager_ImportV2_Call) RunAndReturn(run func(int64, *datapb.ImportRequest) error) *MockDataNodeManager_ImportV2_Call {
_c.Call.Return(run)
return _c
}
// NotifyChannelOperation provides a mock function with given fields: ctx, nodeID, req
func (_m *MockDataNodeManager) NotifyChannelOperation(ctx context.Context, nodeID int64, req *datapb.ChannelOperationsRequest) error {
ret := _m.Called(ctx, nodeID, req)
if len(ret) == 0 {
panic("no return value specified for NotifyChannelOperation")
}
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, int64, *datapb.ChannelOperationsRequest) error); ok {
r0 = rf(ctx, nodeID, req)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockDataNodeManager_NotifyChannelOperation_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'NotifyChannelOperation'
type MockDataNodeManager_NotifyChannelOperation_Call struct {
*mock.Call
}
// NotifyChannelOperation is a helper method to define mock.On call
// - ctx context.Context
// - nodeID int64
// - req *datapb.ChannelOperationsRequest
func (_e *MockDataNodeManager_Expecter) NotifyChannelOperation(ctx interface{}, nodeID interface{}, req interface{}) *MockDataNodeManager_NotifyChannelOperation_Call {
return &MockDataNodeManager_NotifyChannelOperation_Call{Call: _e.mock.On("NotifyChannelOperation", ctx, nodeID, req)}
}
func (_c *MockDataNodeManager_NotifyChannelOperation_Call) Run(run func(ctx context.Context, nodeID int64, req *datapb.ChannelOperationsRequest)) *MockDataNodeManager_NotifyChannelOperation_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(int64), args[2].(*datapb.ChannelOperationsRequest))
})
return _c
}
func (_c *MockDataNodeManager_NotifyChannelOperation_Call) Return(_a0 error) *MockDataNodeManager_NotifyChannelOperation_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockDataNodeManager_NotifyChannelOperation_Call) RunAndReturn(run func(context.Context, int64, *datapb.ChannelOperationsRequest) error) *MockDataNodeManager_NotifyChannelOperation_Call {
_c.Call.Return(run)
return _c
}
// PreImport provides a mock function with given fields: nodeID, in
func (_m *MockDataNodeManager) PreImport(nodeID int64, in *datapb.PreImportRequest) error {
ret := _m.Called(nodeID, in)
if len(ret) == 0 {
panic("no return value specified for PreImport")
}
var r0 error
if rf, ok := ret.Get(0).(func(int64, *datapb.PreImportRequest) error); ok {
r0 = rf(nodeID, in)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockDataNodeManager_PreImport_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'PreImport'
type MockDataNodeManager_PreImport_Call struct {
*mock.Call
}
// PreImport is a helper method to define mock.On call
// - nodeID int64
// - in *datapb.PreImportRequest
func (_e *MockDataNodeManager_Expecter) PreImport(nodeID interface{}, in interface{}) *MockDataNodeManager_PreImport_Call {
return &MockDataNodeManager_PreImport_Call{Call: _e.mock.On("PreImport", nodeID, in)}
}
func (_c *MockDataNodeManager_PreImport_Call) Run(run func(nodeID int64, in *datapb.PreImportRequest)) *MockDataNodeManager_PreImport_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64), args[1].(*datapb.PreImportRequest))
})
return _c
}
func (_c *MockDataNodeManager_PreImport_Call) Return(_a0 error) *MockDataNodeManager_PreImport_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockDataNodeManager_PreImport_Call) RunAndReturn(run func(int64, *datapb.PreImportRequest) error) *MockDataNodeManager_PreImport_Call {
_c.Call.Return(run)
return _c
}
// QueryImport provides a mock function with given fields: nodeID, in
func (_m *MockDataNodeManager) 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
}
// MockDataNodeManager_QueryImport_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'QueryImport'
type MockDataNodeManager_QueryImport_Call struct {
*mock.Call
}
// QueryImport is a helper method to define mock.On call
// - nodeID int64
// - in *datapb.QueryImportRequest
func (_e *MockDataNodeManager_Expecter) QueryImport(nodeID interface{}, in interface{}) *MockDataNodeManager_QueryImport_Call {
return &MockDataNodeManager_QueryImport_Call{Call: _e.mock.On("QueryImport", nodeID, in)}
}
func (_c *MockDataNodeManager_QueryImport_Call) Run(run func(nodeID int64, in *datapb.QueryImportRequest)) *MockDataNodeManager_QueryImport_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64), args[1].(*datapb.QueryImportRequest))
})
return _c
}
func (_c *MockDataNodeManager_QueryImport_Call) Return(_a0 *datapb.QueryImportResponse, _a1 error) *MockDataNodeManager_QueryImport_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockDataNodeManager_QueryImport_Call) RunAndReturn(run func(int64, *datapb.QueryImportRequest) (*datapb.QueryImportResponse, error)) *MockDataNodeManager_QueryImport_Call {
_c.Call.Return(run)
return _c
}
// QueryPreImport provides a mock function with given fields: nodeID, in
func (_m *MockDataNodeManager) 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
}
// MockDataNodeManager_QueryPreImport_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'QueryPreImport'
type MockDataNodeManager_QueryPreImport_Call struct {
*mock.Call
}
// QueryPreImport is a helper method to define mock.On call
// - nodeID int64
// - in *datapb.QueryPreImportRequest
func (_e *MockDataNodeManager_Expecter) QueryPreImport(nodeID interface{}, in interface{}) *MockDataNodeManager_QueryPreImport_Call {
return &MockDataNodeManager_QueryPreImport_Call{Call: _e.mock.On("QueryPreImport", nodeID, in)}
}
func (_c *MockDataNodeManager_QueryPreImport_Call) Run(run func(nodeID int64, in *datapb.QueryPreImportRequest)) *MockDataNodeManager_QueryPreImport_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64), args[1].(*datapb.QueryPreImportRequest))
})
return _c
}
func (_c *MockDataNodeManager_QueryPreImport_Call) Return(_a0 *datapb.QueryPreImportResponse, _a1 error) *MockDataNodeManager_QueryPreImport_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockDataNodeManager_QueryPreImport_Call) RunAndReturn(run func(int64, *datapb.QueryPreImportRequest) (*datapb.QueryPreImportResponse, error)) *MockDataNodeManager_QueryPreImport_Call {
_c.Call.Return(run)
return _c
}
// QuerySlot provides a mock function with given fields: nodeID
func (_m *MockDataNodeManager) QuerySlot(nodeID int64) (*datapb.QuerySlotResponse, error) {
ret := _m.Called(nodeID)
if len(ret) == 0 {
panic("no return value specified for QuerySlot")
}
var r0 *datapb.QuerySlotResponse
var r1 error
if rf, ok := ret.Get(0).(func(int64) (*datapb.QuerySlotResponse, error)); ok {
return rf(nodeID)
}
if rf, ok := ret.Get(0).(func(int64) *datapb.QuerySlotResponse); ok {
r0 = rf(nodeID)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(*datapb.QuerySlotResponse)
}
}
if rf, ok := ret.Get(1).(func(int64) error); ok {
r1 = rf(nodeID)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MockDataNodeManager_QuerySlot_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'QuerySlot'
type MockDataNodeManager_QuerySlot_Call struct {
*mock.Call
}
// QuerySlot is a helper method to define mock.On call
// - nodeID int64
func (_e *MockDataNodeManager_Expecter) QuerySlot(nodeID interface{}) *MockDataNodeManager_QuerySlot_Call {
return &MockDataNodeManager_QuerySlot_Call{Call: _e.mock.On("QuerySlot", nodeID)}
}
func (_c *MockDataNodeManager_QuerySlot_Call) Run(run func(nodeID int64)) *MockDataNodeManager_QuerySlot_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(int64))
})
return _c
}
func (_c *MockDataNodeManager_QuerySlot_Call) Return(_a0 *datapb.QuerySlotResponse, _a1 error) *MockDataNodeManager_QuerySlot_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockDataNodeManager_QuerySlot_Call) RunAndReturn(run func(int64) (*datapb.QuerySlotResponse, error)) *MockDataNodeManager_QuerySlot_Call {
_c.Call.Return(run)
return _c
}
// SyncSegments provides a mock function with given fields: ctx, nodeID, req
func (_m *MockDataNodeManager) SyncSegments(ctx context.Context, nodeID int64, req *datapb.SyncSegmentsRequest) error {
ret := _m.Called(ctx, nodeID, req)
if len(ret) == 0 {
panic("no return value specified for SyncSegments")
}
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, int64, *datapb.SyncSegmentsRequest) error); ok {
r0 = rf(ctx, nodeID, req)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockDataNodeManager_SyncSegments_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SyncSegments'
type MockDataNodeManager_SyncSegments_Call struct {
*mock.Call
}
// SyncSegments is a helper method to define mock.On call
// - ctx context.Context
// - nodeID int64
// - req *datapb.SyncSegmentsRequest
func (_e *MockDataNodeManager_Expecter) SyncSegments(ctx interface{}, nodeID interface{}, req interface{}) *MockDataNodeManager_SyncSegments_Call {
return &MockDataNodeManager_SyncSegments_Call{Call: _e.mock.On("SyncSegments", ctx, nodeID, req)}
}
func (_c *MockDataNodeManager_SyncSegments_Call) Run(run func(ctx context.Context, nodeID int64, req *datapb.SyncSegmentsRequest)) *MockDataNodeManager_SyncSegments_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(int64), args[2].(*datapb.SyncSegmentsRequest))
})
return _c
}
func (_c *MockDataNodeManager_SyncSegments_Call) Return(_a0 error) *MockDataNodeManager_SyncSegments_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockDataNodeManager_SyncSegments_Call) RunAndReturn(run func(context.Context, int64, *datapb.SyncSegmentsRequest) error) *MockDataNodeManager_SyncSegments_Call {
_c.Call.Return(run)
return _c
}
// NewMockDataNodeManager creates a new instance of MockDataNodeManager. 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 NewMockDataNodeManager(t interface {
mock.TestingT
Cleanup(func())
}) *MockDataNodeManager {
mock := &MockDataNodeManager{}
mock.Mock.Test(t)
t.Cleanup(func() { mock.AssertExpectations(t) })
return mock
}