mirror of https://github.com/milvus-io/milvus.git
424 lines
14 KiB
Go
424 lines
14 KiB
Go
// Code generated by mockery v2.53.3. DO NOT EDIT.
|
|
|
|
package mock_balancer
|
|
|
|
import (
|
|
context "context"
|
|
|
|
balancer "github.com/milvus-io/milvus/internal/streamingcoord/server/balancer"
|
|
|
|
mock "github.com/stretchr/testify/mock"
|
|
|
|
streamingpb "github.com/milvus-io/milvus/pkg/v2/proto/streamingpb"
|
|
|
|
syncutil "github.com/milvus-io/milvus/pkg/v2/util/syncutil"
|
|
|
|
types "github.com/milvus-io/milvus/pkg/v2/streaming/util/types"
|
|
)
|
|
|
|
// MockBalancer is an autogenerated mock type for the Balancer type
|
|
type MockBalancer struct {
|
|
mock.Mock
|
|
}
|
|
|
|
type MockBalancer_Expecter struct {
|
|
mock *mock.Mock
|
|
}
|
|
|
|
func (_m *MockBalancer) EXPECT() *MockBalancer_Expecter {
|
|
return &MockBalancer_Expecter{mock: &_m.Mock}
|
|
}
|
|
|
|
// Close provides a mock function with no fields
|
|
func (_m *MockBalancer) Close() {
|
|
_m.Called()
|
|
}
|
|
|
|
// MockBalancer_Close_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Close'
|
|
type MockBalancer_Close_Call struct {
|
|
*mock.Call
|
|
}
|
|
|
|
// Close is a helper method to define mock.On call
|
|
func (_e *MockBalancer_Expecter) Close() *MockBalancer_Close_Call {
|
|
return &MockBalancer_Close_Call{Call: _e.mock.On("Close")}
|
|
}
|
|
|
|
func (_c *MockBalancer_Close_Call) Run(run func()) *MockBalancer_Close_Call {
|
|
_c.Call.Run(func(args mock.Arguments) {
|
|
run()
|
|
})
|
|
return _c
|
|
}
|
|
|
|
func (_c *MockBalancer_Close_Call) Return() *MockBalancer_Close_Call {
|
|
_c.Call.Return()
|
|
return _c
|
|
}
|
|
|
|
func (_c *MockBalancer_Close_Call) RunAndReturn(run func()) *MockBalancer_Close_Call {
|
|
_c.Run(run)
|
|
return _c
|
|
}
|
|
|
|
// GetAllStreamingNodes provides a mock function with given fields: ctx
|
|
func (_m *MockBalancer) GetAllStreamingNodes(ctx context.Context) (map[int64]*types.StreamingNodeInfo, error) {
|
|
ret := _m.Called(ctx)
|
|
|
|
if len(ret) == 0 {
|
|
panic("no return value specified for GetAllStreamingNodes")
|
|
}
|
|
|
|
var r0 map[int64]*types.StreamingNodeInfo
|
|
var r1 error
|
|
if rf, ok := ret.Get(0).(func(context.Context) (map[int64]*types.StreamingNodeInfo, error)); ok {
|
|
return rf(ctx)
|
|
}
|
|
if rf, ok := ret.Get(0).(func(context.Context) map[int64]*types.StreamingNodeInfo); ok {
|
|
r0 = rf(ctx)
|
|
} else {
|
|
if ret.Get(0) != nil {
|
|
r0 = ret.Get(0).(map[int64]*types.StreamingNodeInfo)
|
|
}
|
|
}
|
|
|
|
if rf, ok := ret.Get(1).(func(context.Context) error); ok {
|
|
r1 = rf(ctx)
|
|
} else {
|
|
r1 = ret.Error(1)
|
|
}
|
|
|
|
return r0, r1
|
|
}
|
|
|
|
// MockBalancer_GetAllStreamingNodes_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetAllStreamingNodes'
|
|
type MockBalancer_GetAllStreamingNodes_Call struct {
|
|
*mock.Call
|
|
}
|
|
|
|
// GetAllStreamingNodes is a helper method to define mock.On call
|
|
// - ctx context.Context
|
|
func (_e *MockBalancer_Expecter) GetAllStreamingNodes(ctx interface{}) *MockBalancer_GetAllStreamingNodes_Call {
|
|
return &MockBalancer_GetAllStreamingNodes_Call{Call: _e.mock.On("GetAllStreamingNodes", ctx)}
|
|
}
|
|
|
|
func (_c *MockBalancer_GetAllStreamingNodes_Call) Run(run func(ctx context.Context)) *MockBalancer_GetAllStreamingNodes_Call {
|
|
_c.Call.Run(func(args mock.Arguments) {
|
|
run(args[0].(context.Context))
|
|
})
|
|
return _c
|
|
}
|
|
|
|
func (_c *MockBalancer_GetAllStreamingNodes_Call) Return(_a0 map[int64]*types.StreamingNodeInfo, _a1 error) *MockBalancer_GetAllStreamingNodes_Call {
|
|
_c.Call.Return(_a0, _a1)
|
|
return _c
|
|
}
|
|
|
|
func (_c *MockBalancer_GetAllStreamingNodes_Call) RunAndReturn(run func(context.Context) (map[int64]*types.StreamingNodeInfo, error)) *MockBalancer_GetAllStreamingNodes_Call {
|
|
_c.Call.Return(run)
|
|
return _c
|
|
}
|
|
|
|
// GetLatestWALLocated provides a mock function with given fields: ctx, pchannel
|
|
func (_m *MockBalancer) GetLatestWALLocated(ctx context.Context, pchannel string) (int64, bool) {
|
|
ret := _m.Called(ctx, pchannel)
|
|
|
|
if len(ret) == 0 {
|
|
panic("no return value specified for GetLatestWALLocated")
|
|
}
|
|
|
|
var r0 int64
|
|
var r1 bool
|
|
if rf, ok := ret.Get(0).(func(context.Context, string) (int64, bool)); ok {
|
|
return rf(ctx, pchannel)
|
|
}
|
|
if rf, ok := ret.Get(0).(func(context.Context, string) int64); ok {
|
|
r0 = rf(ctx, pchannel)
|
|
} else {
|
|
r0 = ret.Get(0).(int64)
|
|
}
|
|
|
|
if rf, ok := ret.Get(1).(func(context.Context, string) bool); ok {
|
|
r1 = rf(ctx, pchannel)
|
|
} else {
|
|
r1 = ret.Get(1).(bool)
|
|
}
|
|
|
|
return r0, r1
|
|
}
|
|
|
|
// MockBalancer_GetLatestWALLocated_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetLatestWALLocated'
|
|
type MockBalancer_GetLatestWALLocated_Call struct {
|
|
*mock.Call
|
|
}
|
|
|
|
// GetLatestWALLocated is a helper method to define mock.On call
|
|
// - ctx context.Context
|
|
// - pchannel string
|
|
func (_e *MockBalancer_Expecter) GetLatestWALLocated(ctx interface{}, pchannel interface{}) *MockBalancer_GetLatestWALLocated_Call {
|
|
return &MockBalancer_GetLatestWALLocated_Call{Call: _e.mock.On("GetLatestWALLocated", ctx, pchannel)}
|
|
}
|
|
|
|
func (_c *MockBalancer_GetLatestWALLocated_Call) Run(run func(ctx context.Context, pchannel string)) *MockBalancer_GetLatestWALLocated_Call {
|
|
_c.Call.Run(func(args mock.Arguments) {
|
|
run(args[0].(context.Context), args[1].(string))
|
|
})
|
|
return _c
|
|
}
|
|
|
|
func (_c *MockBalancer_GetLatestWALLocated_Call) Return(_a0 int64, _a1 bool) *MockBalancer_GetLatestWALLocated_Call {
|
|
_c.Call.Return(_a0, _a1)
|
|
return _c
|
|
}
|
|
|
|
func (_c *MockBalancer_GetLatestWALLocated_Call) RunAndReturn(run func(context.Context, string) (int64, bool)) *MockBalancer_GetLatestWALLocated_Call {
|
|
_c.Call.Return(run)
|
|
return _c
|
|
}
|
|
|
|
// MarkAsUnavailable provides a mock function with given fields: ctx, pChannels
|
|
func (_m *MockBalancer) MarkAsUnavailable(ctx context.Context, pChannels []types.PChannelInfo) error {
|
|
ret := _m.Called(ctx, pChannels)
|
|
|
|
if len(ret) == 0 {
|
|
panic("no return value specified for MarkAsUnavailable")
|
|
}
|
|
|
|
var r0 error
|
|
if rf, ok := ret.Get(0).(func(context.Context, []types.PChannelInfo) error); ok {
|
|
r0 = rf(ctx, pChannels)
|
|
} else {
|
|
r0 = ret.Error(0)
|
|
}
|
|
|
|
return r0
|
|
}
|
|
|
|
// MockBalancer_MarkAsUnavailable_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'MarkAsUnavailable'
|
|
type MockBalancer_MarkAsUnavailable_Call struct {
|
|
*mock.Call
|
|
}
|
|
|
|
// MarkAsUnavailable is a helper method to define mock.On call
|
|
// - ctx context.Context
|
|
// - pChannels []types.PChannelInfo
|
|
func (_e *MockBalancer_Expecter) MarkAsUnavailable(ctx interface{}, pChannels interface{}) *MockBalancer_MarkAsUnavailable_Call {
|
|
return &MockBalancer_MarkAsUnavailable_Call{Call: _e.mock.On("MarkAsUnavailable", ctx, pChannels)}
|
|
}
|
|
|
|
func (_c *MockBalancer_MarkAsUnavailable_Call) Run(run func(ctx context.Context, pChannels []types.PChannelInfo)) *MockBalancer_MarkAsUnavailable_Call {
|
|
_c.Call.Run(func(args mock.Arguments) {
|
|
run(args[0].(context.Context), args[1].([]types.PChannelInfo))
|
|
})
|
|
return _c
|
|
}
|
|
|
|
func (_c *MockBalancer_MarkAsUnavailable_Call) Return(_a0 error) *MockBalancer_MarkAsUnavailable_Call {
|
|
_c.Call.Return(_a0)
|
|
return _c
|
|
}
|
|
|
|
func (_c *MockBalancer_MarkAsUnavailable_Call) RunAndReturn(run func(context.Context, []types.PChannelInfo) error) *MockBalancer_MarkAsUnavailable_Call {
|
|
_c.Call.Return(run)
|
|
return _c
|
|
}
|
|
|
|
// RegisterStreamingEnabledNotifier provides a mock function with given fields: notifier
|
|
func (_m *MockBalancer) RegisterStreamingEnabledNotifier(notifier *syncutil.AsyncTaskNotifier[struct{}]) {
|
|
_m.Called(notifier)
|
|
}
|
|
|
|
// MockBalancer_RegisterStreamingEnabledNotifier_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RegisterStreamingEnabledNotifier'
|
|
type MockBalancer_RegisterStreamingEnabledNotifier_Call struct {
|
|
*mock.Call
|
|
}
|
|
|
|
// RegisterStreamingEnabledNotifier is a helper method to define mock.On call
|
|
// - notifier *syncutil.AsyncTaskNotifier[struct{}]
|
|
func (_e *MockBalancer_Expecter) RegisterStreamingEnabledNotifier(notifier interface{}) *MockBalancer_RegisterStreamingEnabledNotifier_Call {
|
|
return &MockBalancer_RegisterStreamingEnabledNotifier_Call{Call: _e.mock.On("RegisterStreamingEnabledNotifier", notifier)}
|
|
}
|
|
|
|
func (_c *MockBalancer_RegisterStreamingEnabledNotifier_Call) Run(run func(notifier *syncutil.AsyncTaskNotifier[struct{}])) *MockBalancer_RegisterStreamingEnabledNotifier_Call {
|
|
_c.Call.Run(func(args mock.Arguments) {
|
|
run(args[0].(*syncutil.AsyncTaskNotifier[struct{}]))
|
|
})
|
|
return _c
|
|
}
|
|
|
|
func (_c *MockBalancer_RegisterStreamingEnabledNotifier_Call) Return() *MockBalancer_RegisterStreamingEnabledNotifier_Call {
|
|
_c.Call.Return()
|
|
return _c
|
|
}
|
|
|
|
func (_c *MockBalancer_RegisterStreamingEnabledNotifier_Call) RunAndReturn(run func(*syncutil.AsyncTaskNotifier[struct{}])) *MockBalancer_RegisterStreamingEnabledNotifier_Call {
|
|
_c.Run(run)
|
|
return _c
|
|
}
|
|
|
|
// Trigger provides a mock function with given fields: ctx
|
|
func (_m *MockBalancer) Trigger(ctx context.Context) error {
|
|
ret := _m.Called(ctx)
|
|
|
|
if len(ret) == 0 {
|
|
panic("no return value specified for Trigger")
|
|
}
|
|
|
|
var r0 error
|
|
if rf, ok := ret.Get(0).(func(context.Context) error); ok {
|
|
r0 = rf(ctx)
|
|
} else {
|
|
r0 = ret.Error(0)
|
|
}
|
|
|
|
return r0
|
|
}
|
|
|
|
// MockBalancer_Trigger_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Trigger'
|
|
type MockBalancer_Trigger_Call struct {
|
|
*mock.Call
|
|
}
|
|
|
|
// Trigger is a helper method to define mock.On call
|
|
// - ctx context.Context
|
|
func (_e *MockBalancer_Expecter) Trigger(ctx interface{}) *MockBalancer_Trigger_Call {
|
|
return &MockBalancer_Trigger_Call{Call: _e.mock.On("Trigger", ctx)}
|
|
}
|
|
|
|
func (_c *MockBalancer_Trigger_Call) Run(run func(ctx context.Context)) *MockBalancer_Trigger_Call {
|
|
_c.Call.Run(func(args mock.Arguments) {
|
|
run(args[0].(context.Context))
|
|
})
|
|
return _c
|
|
}
|
|
|
|
func (_c *MockBalancer_Trigger_Call) Return(_a0 error) *MockBalancer_Trigger_Call {
|
|
_c.Call.Return(_a0)
|
|
return _c
|
|
}
|
|
|
|
func (_c *MockBalancer_Trigger_Call) RunAndReturn(run func(context.Context) error) *MockBalancer_Trigger_Call {
|
|
_c.Call.Return(run)
|
|
return _c
|
|
}
|
|
|
|
// UpdateBalancePolicy provides a mock function with given fields: ctx, req
|
|
func (_m *MockBalancer) UpdateBalancePolicy(ctx context.Context, req *streamingpb.UpdateWALBalancePolicyRequest) (*streamingpb.UpdateWALBalancePolicyResponse, error) {
|
|
ret := _m.Called(ctx, req)
|
|
|
|
if len(ret) == 0 {
|
|
panic("no return value specified for UpdateBalancePolicy")
|
|
}
|
|
|
|
var r0 *streamingpb.UpdateWALBalancePolicyResponse
|
|
var r1 error
|
|
if rf, ok := ret.Get(0).(func(context.Context, *streamingpb.UpdateWALBalancePolicyRequest) (*streamingpb.UpdateWALBalancePolicyResponse, error)); ok {
|
|
return rf(ctx, req)
|
|
}
|
|
if rf, ok := ret.Get(0).(func(context.Context, *streamingpb.UpdateWALBalancePolicyRequest) *streamingpb.UpdateWALBalancePolicyResponse); ok {
|
|
r0 = rf(ctx, req)
|
|
} else {
|
|
if ret.Get(0) != nil {
|
|
r0 = ret.Get(0).(*streamingpb.UpdateWALBalancePolicyResponse)
|
|
}
|
|
}
|
|
|
|
if rf, ok := ret.Get(1).(func(context.Context, *streamingpb.UpdateWALBalancePolicyRequest) error); ok {
|
|
r1 = rf(ctx, req)
|
|
} else {
|
|
r1 = ret.Error(1)
|
|
}
|
|
|
|
return r0, r1
|
|
}
|
|
|
|
// MockBalancer_UpdateBalancePolicy_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'UpdateBalancePolicy'
|
|
type MockBalancer_UpdateBalancePolicy_Call struct {
|
|
*mock.Call
|
|
}
|
|
|
|
// UpdateBalancePolicy is a helper method to define mock.On call
|
|
// - ctx context.Context
|
|
// - req *streamingpb.UpdateWALBalancePolicyRequest
|
|
func (_e *MockBalancer_Expecter) UpdateBalancePolicy(ctx interface{}, req interface{}) *MockBalancer_UpdateBalancePolicy_Call {
|
|
return &MockBalancer_UpdateBalancePolicy_Call{Call: _e.mock.On("UpdateBalancePolicy", ctx, req)}
|
|
}
|
|
|
|
func (_c *MockBalancer_UpdateBalancePolicy_Call) Run(run func(ctx context.Context, req *streamingpb.UpdateWALBalancePolicyRequest)) *MockBalancer_UpdateBalancePolicy_Call {
|
|
_c.Call.Run(func(args mock.Arguments) {
|
|
run(args[0].(context.Context), args[1].(*streamingpb.UpdateWALBalancePolicyRequest))
|
|
})
|
|
return _c
|
|
}
|
|
|
|
func (_c *MockBalancer_UpdateBalancePolicy_Call) Return(_a0 *streamingpb.UpdateWALBalancePolicyResponse, _a1 error) *MockBalancer_UpdateBalancePolicy_Call {
|
|
_c.Call.Return(_a0, _a1)
|
|
return _c
|
|
}
|
|
|
|
func (_c *MockBalancer_UpdateBalancePolicy_Call) RunAndReturn(run func(context.Context, *streamingpb.UpdateWALBalancePolicyRequest) (*streamingpb.UpdateWALBalancePolicyResponse, error)) *MockBalancer_UpdateBalancePolicy_Call {
|
|
_c.Call.Return(run)
|
|
return _c
|
|
}
|
|
|
|
// WatchChannelAssignments provides a mock function with given fields: ctx, cb
|
|
func (_m *MockBalancer) WatchChannelAssignments(ctx context.Context, cb balancer.WatchChannelAssignmentsCallback) error {
|
|
ret := _m.Called(ctx, cb)
|
|
|
|
if len(ret) == 0 {
|
|
panic("no return value specified for WatchChannelAssignments")
|
|
}
|
|
|
|
var r0 error
|
|
if rf, ok := ret.Get(0).(func(context.Context, balancer.WatchChannelAssignmentsCallback) error); ok {
|
|
r0 = rf(ctx, cb)
|
|
} else {
|
|
r0 = ret.Error(0)
|
|
}
|
|
|
|
return r0
|
|
}
|
|
|
|
// MockBalancer_WatchChannelAssignments_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'WatchChannelAssignments'
|
|
type MockBalancer_WatchChannelAssignments_Call struct {
|
|
*mock.Call
|
|
}
|
|
|
|
// WatchChannelAssignments is a helper method to define mock.On call
|
|
// - ctx context.Context
|
|
// - cb balancer.WatchChannelAssignmentsCallback
|
|
func (_e *MockBalancer_Expecter) WatchChannelAssignments(ctx interface{}, cb interface{}) *MockBalancer_WatchChannelAssignments_Call {
|
|
return &MockBalancer_WatchChannelAssignments_Call{Call: _e.mock.On("WatchChannelAssignments", ctx, cb)}
|
|
}
|
|
|
|
func (_c *MockBalancer_WatchChannelAssignments_Call) Run(run func(ctx context.Context, cb balancer.WatchChannelAssignmentsCallback)) *MockBalancer_WatchChannelAssignments_Call {
|
|
_c.Call.Run(func(args mock.Arguments) {
|
|
run(args[0].(context.Context), args[1].(balancer.WatchChannelAssignmentsCallback))
|
|
})
|
|
return _c
|
|
}
|
|
|
|
func (_c *MockBalancer_WatchChannelAssignments_Call) Return(_a0 error) *MockBalancer_WatchChannelAssignments_Call {
|
|
_c.Call.Return(_a0)
|
|
return _c
|
|
}
|
|
|
|
func (_c *MockBalancer_WatchChannelAssignments_Call) RunAndReturn(run func(context.Context, balancer.WatchChannelAssignmentsCallback) error) *MockBalancer_WatchChannelAssignments_Call {
|
|
_c.Call.Return(run)
|
|
return _c
|
|
}
|
|
|
|
// NewMockBalancer creates a new instance of MockBalancer. 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 NewMockBalancer(t interface {
|
|
mock.TestingT
|
|
Cleanup(func())
|
|
}) *MockBalancer {
|
|
mock := &MockBalancer{}
|
|
mock.Mock.Test(t)
|
|
|
|
t.Cleanup(func() { mock.AssertExpectations(t) })
|
|
|
|
return mock
|
|
}
|