milvus/internal/datanode/channel/mock_channel_manager.go
yihao.dai 36e9e41627
fix: Fix no candidate segments error for small import (#41771)
When autoID is enabled, the preimport task estimates row distribution by
evenly dividing the total row count (numRows) across all vchannels:
`estimatedCount = numRows / vchannelNum`.
However, the actual import task hashes real auto-generated IDs to
determine
the target vchannel. This mismatch can lead to inaccurate row
distribution estimation
in such corner cases:
- Importing 1 row into 2 vchannels:
				• Preimport: 1 / 2 = 0 → both v0 and v1 are estimated to have 0 rows
				• Import: real autoID (e.g., 457975852966809057) hashes to v1
				  → actual result: v0 = 0, v1 = 1

To resolve such corner case, we now allocate at least one segment for
each vchannel
when autoID is enabled, ensuring all vchannels are prepared to receive
data even
if no rows are estimated for them.

issue: https://github.com/milvus-io/milvus/issues/41759

---------

Signed-off-by: bigsheeper <yihao.dai@zilliz.com>
2025-05-14 15:30:21 +08:00

194 lines
5.6 KiB
Go

// Code generated by mockery v2.53.3. DO NOT EDIT.
package channel
import (
datapb "github.com/milvus-io/milvus/pkg/v2/proto/datapb"
mock "github.com/stretchr/testify/mock"
)
// MockChannelManager is an autogenerated mock type for the ChannelManager type
type MockChannelManager struct {
mock.Mock
}
type MockChannelManager_Expecter struct {
mock *mock.Mock
}
func (_m *MockChannelManager) EXPECT() *MockChannelManager_Expecter {
return &MockChannelManager_Expecter{mock: &_m.Mock}
}
// Close provides a mock function with no fields
func (_m *MockChannelManager) Close() {
_m.Called()
}
// MockChannelManager_Close_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Close'
type MockChannelManager_Close_Call struct {
*mock.Call
}
// Close is a helper method to define mock.On call
func (_e *MockChannelManager_Expecter) Close() *MockChannelManager_Close_Call {
return &MockChannelManager_Close_Call{Call: _e.mock.On("Close")}
}
func (_c *MockChannelManager_Close_Call) Run(run func()) *MockChannelManager_Close_Call {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *MockChannelManager_Close_Call) Return() *MockChannelManager_Close_Call {
_c.Call.Return()
return _c
}
func (_c *MockChannelManager_Close_Call) RunAndReturn(run func()) *MockChannelManager_Close_Call {
_c.Run(run)
return _c
}
// GetProgress provides a mock function with given fields: info
func (_m *MockChannelManager) GetProgress(info *datapb.ChannelWatchInfo) *datapb.ChannelOperationProgressResponse {
ret := _m.Called(info)
if len(ret) == 0 {
panic("no return value specified for GetProgress")
}
var r0 *datapb.ChannelOperationProgressResponse
if rf, ok := ret.Get(0).(func(*datapb.ChannelWatchInfo) *datapb.ChannelOperationProgressResponse); ok {
r0 = rf(info)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(*datapb.ChannelOperationProgressResponse)
}
}
return r0
}
// MockChannelManager_GetProgress_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetProgress'
type MockChannelManager_GetProgress_Call struct {
*mock.Call
}
// GetProgress is a helper method to define mock.On call
// - info *datapb.ChannelWatchInfo
func (_e *MockChannelManager_Expecter) GetProgress(info interface{}) *MockChannelManager_GetProgress_Call {
return &MockChannelManager_GetProgress_Call{Call: _e.mock.On("GetProgress", info)}
}
func (_c *MockChannelManager_GetProgress_Call) Run(run func(info *datapb.ChannelWatchInfo)) *MockChannelManager_GetProgress_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(*datapb.ChannelWatchInfo))
})
return _c
}
func (_c *MockChannelManager_GetProgress_Call) Return(_a0 *datapb.ChannelOperationProgressResponse) *MockChannelManager_GetProgress_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockChannelManager_GetProgress_Call) RunAndReturn(run func(*datapb.ChannelWatchInfo) *datapb.ChannelOperationProgressResponse) *MockChannelManager_GetProgress_Call {
_c.Call.Return(run)
return _c
}
// Start provides a mock function with no fields
func (_m *MockChannelManager) Start() {
_m.Called()
}
// MockChannelManager_Start_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Start'
type MockChannelManager_Start_Call struct {
*mock.Call
}
// Start is a helper method to define mock.On call
func (_e *MockChannelManager_Expecter) Start() *MockChannelManager_Start_Call {
return &MockChannelManager_Start_Call{Call: _e.mock.On("Start")}
}
func (_c *MockChannelManager_Start_Call) Run(run func()) *MockChannelManager_Start_Call {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *MockChannelManager_Start_Call) Return() *MockChannelManager_Start_Call {
_c.Call.Return()
return _c
}
func (_c *MockChannelManager_Start_Call) RunAndReturn(run func()) *MockChannelManager_Start_Call {
_c.Run(run)
return _c
}
// Submit provides a mock function with given fields: info
func (_m *MockChannelManager) Submit(info *datapb.ChannelWatchInfo) error {
ret := _m.Called(info)
if len(ret) == 0 {
panic("no return value specified for Submit")
}
var r0 error
if rf, ok := ret.Get(0).(func(*datapb.ChannelWatchInfo) error); ok {
r0 = rf(info)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockChannelManager_Submit_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Submit'
type MockChannelManager_Submit_Call struct {
*mock.Call
}
// Submit is a helper method to define mock.On call
// - info *datapb.ChannelWatchInfo
func (_e *MockChannelManager_Expecter) Submit(info interface{}) *MockChannelManager_Submit_Call {
return &MockChannelManager_Submit_Call{Call: _e.mock.On("Submit", info)}
}
func (_c *MockChannelManager_Submit_Call) Run(run func(info *datapb.ChannelWatchInfo)) *MockChannelManager_Submit_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(*datapb.ChannelWatchInfo))
})
return _c
}
func (_c *MockChannelManager_Submit_Call) Return(_a0 error) *MockChannelManager_Submit_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockChannelManager_Submit_Call) RunAndReturn(run func(*datapb.ChannelWatchInfo) error) *MockChannelManager_Submit_Call {
_c.Call.Return(run)
return _c
}
// NewMockChannelManager creates a new instance of MockChannelManager. 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 NewMockChannelManager(t interface {
mock.TestingT
Cleanup(func())
}) *MockChannelManager {
mock := &MockChannelManager{}
mock.Mock.Test(t)
t.Cleanup(func() { mock.AssertExpectations(t) })
return mock
}