milvus/internal/mocks/streamingcoord/server/mock_balancer/mock_Balancer.go
chyezh fda720b880
enhance: streaming service grpc utilities (#34436)
issue: #33285

- add two grpc resolver (by session and by streaming coord assignment
service)
- add one grpc balancer (by serverID and roundrobin)
- add lazy conn to avoid block by first service discovery
- add some utility function for streaming service

Signed-off-by: chyezh <chyezh@outlook.com>
2024-07-15 20:49:38 +08:00

200 lines
6.3 KiB
Go

// Code generated by mockery v2.32.4. DO NOT EDIT.
package mock_balancer
import (
context "context"
types "github.com/milvus-io/milvus/pkg/streaming/util/types"
mock "github.com/stretchr/testify/mock"
typeutil "github.com/milvus-io/milvus/pkg/util/typeutil"
)
// 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 given 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.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)
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
}
// Trigger provides a mock function with given fields: ctx
func (_m *MockBalancer) Trigger(ctx context.Context) error {
ret := _m.Called(ctx)
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
}
// WatchChannelAssignments provides a mock function with given fields: ctx, cb
func (_m *MockBalancer) WatchChannelAssignments(ctx context.Context, cb func(typeutil.VersionInt64Pair, []types.PChannelInfoAssigned) error) error {
ret := _m.Called(ctx, cb)
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, func(typeutil.VersionInt64Pair, []types.PChannelInfoAssigned) error) 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 func(typeutil.VersionInt64Pair , []types.PChannelInfoAssigned) error
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 func(typeutil.VersionInt64Pair, []types.PChannelInfoAssigned) error)) *MockBalancer_WatchChannelAssignments_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(func(typeutil.VersionInt64Pair, []types.PChannelInfoAssigned) error))
})
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, func(typeutil.VersionInt64Pair, []types.PChannelInfoAssigned) error) 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
}