From d19a82e3f6bc59c8603d8bc21f1d1e0e625a12d5 Mon Sep 17 00:00:00 2001 From: aoiasd <45024769+aoiasd@users.noreply.github.com> Date: Thu, 7 Mar 2024 17:45:01 +0800 Subject: [PATCH] enhance: support stream call for grpc client (#30013) relate: https://github.com/milvus-io/milvus/issues/29719 Signed-off-by: aoiasd --- internal/util/grpcclient/client.go | 3 + internal/util/grpcclient/client_test.go | 9 + .../util/streamrpc/mock_grpc_client_stream.go | 302 ++++++++++++++++ .../mock_query_stream_segments_server.go | 325 ------------------ .../mocks/mock_query_stream_server.go | 325 ------------------ 5 files changed, 314 insertions(+), 650 deletions(-) create mode 100644 internal/util/streamrpc/mock_grpc_client_stream.go delete mode 100644 internal/util/streamrpc/mocks/mock_query_stream_segments_server.go delete mode 100644 internal/util/streamrpc/mocks/mock_query_stream_server.go diff --git a/internal/util/grpcclient/client.go b/internal/util/grpcclient/client.go index 0ac7137ce6..d5cf30fdcb 100644 --- a/internal/util/grpcclient/client.go +++ b/internal/util/grpcclient/client.go @@ -509,6 +509,9 @@ func (c *ClientBase[T]) call(ctx context.Context, caller func(client T) (any, er status = res case interface{ GetStatus() *commonpb.Status }: status = res.GetStatus() + // streaming call + case grpc.ClientStream: + status = merr.Status(nil) default: // it will directly return the result log.Warn("unknown return type", zap.Any("return", ret)) diff --git a/internal/util/grpcclient/client_test.go b/internal/util/grpcclient/client_test.go index 8aa8aad139..e5f9078562 100644 --- a/internal/util/grpcclient/client_test.go +++ b/internal/util/grpcclient/client_test.go @@ -40,6 +40,7 @@ import ( "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb" "github.com/milvus-io/milvus/internal/proto/rootcoordpb" "github.com/milvus-io/milvus/internal/util/sessionutil" + "github.com/milvus-io/milvus/internal/util/streamrpc" "github.com/milvus-io/milvus/pkg/util/merr" "github.com/milvus-io/milvus/pkg/util/paramtable" "github.com/milvus-io/milvus/pkg/util/typeutil" @@ -171,6 +172,14 @@ func testCall(t *testing.T, compressed bool) { assert.NoError(t, err) }) + t.Run("Call with stream method", func(t *testing.T) { + initClient() + _, err := base.Call(context.Background(), func(client *mockClient) (any, error) { + return streamrpc.NewMockClientStream(t), nil + }) + assert.NoError(t, err) + }) + t.Run("Call with canceled context", func(t *testing.T) { initClient() ctx, cancel := context.WithCancel(context.Background()) diff --git a/internal/util/streamrpc/mock_grpc_client_stream.go b/internal/util/streamrpc/mock_grpc_client_stream.go new file mode 100644 index 0000000000..8985153b26 --- /dev/null +++ b/internal/util/streamrpc/mock_grpc_client_stream.go @@ -0,0 +1,302 @@ +// Code generated by mockery v2.32.4. DO NOT EDIT. + +package streamrpc + +import ( + context "context" + + metadata "google.golang.org/grpc/metadata" + + mock "github.com/stretchr/testify/mock" +) + +// MockClientStream is an autogenerated mock type for the ClientStream type +type MockClientStream struct { + mock.Mock +} + +type MockClientStream_Expecter struct { + mock *mock.Mock +} + +func (_m *MockClientStream) EXPECT() *MockClientStream_Expecter { + return &MockClientStream_Expecter{mock: &_m.Mock} +} + +// CloseSend provides a mock function with given fields: +func (_m *MockClientStream) CloseSend() error { + ret := _m.Called() + + var r0 error + if rf, ok := ret.Get(0).(func() error); ok { + r0 = rf() + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// MockClientStream_CloseSend_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CloseSend' +type MockClientStream_CloseSend_Call struct { + *mock.Call +} + +// CloseSend is a helper method to define mock.On call +func (_e *MockClientStream_Expecter) CloseSend() *MockClientStream_CloseSend_Call { + return &MockClientStream_CloseSend_Call{Call: _e.mock.On("CloseSend")} +} + +func (_c *MockClientStream_CloseSend_Call) Run(run func()) *MockClientStream_CloseSend_Call { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *MockClientStream_CloseSend_Call) Return(_a0 error) *MockClientStream_CloseSend_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockClientStream_CloseSend_Call) RunAndReturn(run func() error) *MockClientStream_CloseSend_Call { + _c.Call.Return(run) + return _c +} + +// Context provides a mock function with given fields: +func (_m *MockClientStream) Context() context.Context { + ret := _m.Called() + + var r0 context.Context + if rf, ok := ret.Get(0).(func() context.Context); ok { + r0 = rf() + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(context.Context) + } + } + + return r0 +} + +// MockClientStream_Context_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Context' +type MockClientStream_Context_Call struct { + *mock.Call +} + +// Context is a helper method to define mock.On call +func (_e *MockClientStream_Expecter) Context() *MockClientStream_Context_Call { + return &MockClientStream_Context_Call{Call: _e.mock.On("Context")} +} + +func (_c *MockClientStream_Context_Call) Run(run func()) *MockClientStream_Context_Call { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *MockClientStream_Context_Call) Return(_a0 context.Context) *MockClientStream_Context_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockClientStream_Context_Call) RunAndReturn(run func() context.Context) *MockClientStream_Context_Call { + _c.Call.Return(run) + return _c +} + +// Header provides a mock function with given fields: +func (_m *MockClientStream) Header() (metadata.MD, error) { + ret := _m.Called() + + var r0 metadata.MD + var r1 error + if rf, ok := ret.Get(0).(func() (metadata.MD, error)); ok { + return rf() + } + if rf, ok := ret.Get(0).(func() metadata.MD); ok { + r0 = rf() + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(metadata.MD) + } + } + + if rf, ok := ret.Get(1).(func() error); ok { + r1 = rf() + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// MockClientStream_Header_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Header' +type MockClientStream_Header_Call struct { + *mock.Call +} + +// Header is a helper method to define mock.On call +func (_e *MockClientStream_Expecter) Header() *MockClientStream_Header_Call { + return &MockClientStream_Header_Call{Call: _e.mock.On("Header")} +} + +func (_c *MockClientStream_Header_Call) Run(run func()) *MockClientStream_Header_Call { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *MockClientStream_Header_Call) Return(_a0 metadata.MD, _a1 error) *MockClientStream_Header_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MockClientStream_Header_Call) RunAndReturn(run func() (metadata.MD, error)) *MockClientStream_Header_Call { + _c.Call.Return(run) + return _c +} + +// RecvMsg provides a mock function with given fields: m +func (_m *MockClientStream) RecvMsg(m interface{}) error { + ret := _m.Called(m) + + var r0 error + if rf, ok := ret.Get(0).(func(interface{}) error); ok { + r0 = rf(m) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// MockClientStream_RecvMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RecvMsg' +type MockClientStream_RecvMsg_Call struct { + *mock.Call +} + +// RecvMsg is a helper method to define mock.On call +// - m interface{} +func (_e *MockClientStream_Expecter) RecvMsg(m interface{}) *MockClientStream_RecvMsg_Call { + return &MockClientStream_RecvMsg_Call{Call: _e.mock.On("RecvMsg", m)} +} + +func (_c *MockClientStream_RecvMsg_Call) Run(run func(m interface{})) *MockClientStream_RecvMsg_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(interface{})) + }) + return _c +} + +func (_c *MockClientStream_RecvMsg_Call) Return(_a0 error) *MockClientStream_RecvMsg_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockClientStream_RecvMsg_Call) RunAndReturn(run func(interface{}) error) *MockClientStream_RecvMsg_Call { + _c.Call.Return(run) + return _c +} + +// SendMsg provides a mock function with given fields: m +func (_m *MockClientStream) SendMsg(m interface{}) error { + ret := _m.Called(m) + + var r0 error + if rf, ok := ret.Get(0).(func(interface{}) error); ok { + r0 = rf(m) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// MockClientStream_SendMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SendMsg' +type MockClientStream_SendMsg_Call struct { + *mock.Call +} + +// SendMsg is a helper method to define mock.On call +// - m interface{} +func (_e *MockClientStream_Expecter) SendMsg(m interface{}) *MockClientStream_SendMsg_Call { + return &MockClientStream_SendMsg_Call{Call: _e.mock.On("SendMsg", m)} +} + +func (_c *MockClientStream_SendMsg_Call) Run(run func(m interface{})) *MockClientStream_SendMsg_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(interface{})) + }) + return _c +} + +func (_c *MockClientStream_SendMsg_Call) Return(_a0 error) *MockClientStream_SendMsg_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockClientStream_SendMsg_Call) RunAndReturn(run func(interface{}) error) *MockClientStream_SendMsg_Call { + _c.Call.Return(run) + return _c +} + +// Trailer provides a mock function with given fields: +func (_m *MockClientStream) Trailer() metadata.MD { + ret := _m.Called() + + var r0 metadata.MD + if rf, ok := ret.Get(0).(func() metadata.MD); ok { + r0 = rf() + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(metadata.MD) + } + } + + return r0 +} + +// MockClientStream_Trailer_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Trailer' +type MockClientStream_Trailer_Call struct { + *mock.Call +} + +// Trailer is a helper method to define mock.On call +func (_e *MockClientStream_Expecter) Trailer() *MockClientStream_Trailer_Call { + return &MockClientStream_Trailer_Call{Call: _e.mock.On("Trailer")} +} + +func (_c *MockClientStream_Trailer_Call) Run(run func()) *MockClientStream_Trailer_Call { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *MockClientStream_Trailer_Call) Return(_a0 metadata.MD) *MockClientStream_Trailer_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockClientStream_Trailer_Call) RunAndReturn(run func() metadata.MD) *MockClientStream_Trailer_Call { + _c.Call.Return(run) + return _c +} + +// NewMockClientStream creates a new instance of MockClientStream. 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 NewMockClientStream(t interface { + mock.TestingT + Cleanup(func()) +}) *MockClientStream { + mock := &MockClientStream{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/internal/util/streamrpc/mocks/mock_query_stream_segments_server.go b/internal/util/streamrpc/mocks/mock_query_stream_segments_server.go deleted file mode 100644 index b71672a998..0000000000 --- a/internal/util/streamrpc/mocks/mock_query_stream_segments_server.go +++ /dev/null @@ -1,325 +0,0 @@ -// Code generated by mockery v2.32.4. DO NOT EDIT. - -package mocks - -import ( - context "context" - - internalpb "github.com/milvus-io/milvus/internal/proto/internalpb" - metadata "google.golang.org/grpc/metadata" - - mock "github.com/stretchr/testify/mock" -) - -// MockQueryStreamSegmentsServer is an autogenerated mock type for the QueryNode_QueryStreamSegmentsServer type -type MockQueryStreamSegmentsServer struct { - mock.Mock -} - -type MockQueryStreamSegmentsServer_Expecter struct { - mock *mock.Mock -} - -func (_m *MockQueryStreamSegmentsServer) EXPECT() *MockQueryStreamSegmentsServer_Expecter { - return &MockQueryStreamSegmentsServer_Expecter{mock: &_m.Mock} -} - -// Context provides a mock function with given fields: -func (_m *MockQueryStreamSegmentsServer) Context() context.Context { - ret := _m.Called() - - var r0 context.Context - if rf, ok := ret.Get(0).(func() context.Context); ok { - r0 = rf() - } else { - if ret.Get(0) != nil { - r0 = ret.Get(0).(context.Context) - } - } - - return r0 -} - -// MockQueryStreamSegmentsServer_Context_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Context' -type MockQueryStreamSegmentsServer_Context_Call struct { - *mock.Call -} - -// Context is a helper method to define mock.On call -func (_e *MockQueryStreamSegmentsServer_Expecter) Context() *MockQueryStreamSegmentsServer_Context_Call { - return &MockQueryStreamSegmentsServer_Context_Call{Call: _e.mock.On("Context")} -} - -func (_c *MockQueryStreamSegmentsServer_Context_Call) Run(run func()) *MockQueryStreamSegmentsServer_Context_Call { - _c.Call.Run(func(args mock.Arguments) { - run() - }) - return _c -} - -func (_c *MockQueryStreamSegmentsServer_Context_Call) Return(_a0 context.Context) *MockQueryStreamSegmentsServer_Context_Call { - _c.Call.Return(_a0) - return _c -} - -func (_c *MockQueryStreamSegmentsServer_Context_Call) RunAndReturn(run func() context.Context) *MockQueryStreamSegmentsServer_Context_Call { - _c.Call.Return(run) - return _c -} - -// RecvMsg provides a mock function with given fields: m -func (_m *MockQueryStreamSegmentsServer) RecvMsg(m interface{}) error { - ret := _m.Called(m) - - var r0 error - if rf, ok := ret.Get(0).(func(interface{}) error); ok { - r0 = rf(m) - } else { - r0 = ret.Error(0) - } - - return r0 -} - -// MockQueryStreamSegmentsServer_RecvMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RecvMsg' -type MockQueryStreamSegmentsServer_RecvMsg_Call struct { - *mock.Call -} - -// RecvMsg is a helper method to define mock.On call -// - m interface{} -func (_e *MockQueryStreamSegmentsServer_Expecter) RecvMsg(m interface{}) *MockQueryStreamSegmentsServer_RecvMsg_Call { - return &MockQueryStreamSegmentsServer_RecvMsg_Call{Call: _e.mock.On("RecvMsg", m)} -} - -func (_c *MockQueryStreamSegmentsServer_RecvMsg_Call) Run(run func(m interface{})) *MockQueryStreamSegmentsServer_RecvMsg_Call { - _c.Call.Run(func(args mock.Arguments) { - run(args[0].(interface{})) - }) - return _c -} - -func (_c *MockQueryStreamSegmentsServer_RecvMsg_Call) Return(_a0 error) *MockQueryStreamSegmentsServer_RecvMsg_Call { - _c.Call.Return(_a0) - return _c -} - -func (_c *MockQueryStreamSegmentsServer_RecvMsg_Call) RunAndReturn(run func(interface{}) error) *MockQueryStreamSegmentsServer_RecvMsg_Call { - _c.Call.Return(run) - return _c -} - -// Send provides a mock function with given fields: _a0 -func (_m *MockQueryStreamSegmentsServer) Send(_a0 *internalpb.RetrieveResults) error { - ret := _m.Called(_a0) - - var r0 error - if rf, ok := ret.Get(0).(func(*internalpb.RetrieveResults) error); ok { - r0 = rf(_a0) - } else { - r0 = ret.Error(0) - } - - return r0 -} - -// MockQueryStreamSegmentsServer_Send_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Send' -type MockQueryStreamSegmentsServer_Send_Call struct { - *mock.Call -} - -// Send is a helper method to define mock.On call -// - _a0 *internalpb.RetrieveResults -func (_e *MockQueryStreamSegmentsServer_Expecter) Send(_a0 interface{}) *MockQueryStreamSegmentsServer_Send_Call { - return &MockQueryStreamSegmentsServer_Send_Call{Call: _e.mock.On("Send", _a0)} -} - -func (_c *MockQueryStreamSegmentsServer_Send_Call) Run(run func(_a0 *internalpb.RetrieveResults)) *MockQueryStreamSegmentsServer_Send_Call { - _c.Call.Run(func(args mock.Arguments) { - run(args[0].(*internalpb.RetrieveResults)) - }) - return _c -} - -func (_c *MockQueryStreamSegmentsServer_Send_Call) Return(_a0 error) *MockQueryStreamSegmentsServer_Send_Call { - _c.Call.Return(_a0) - return _c -} - -func (_c *MockQueryStreamSegmentsServer_Send_Call) RunAndReturn(run func(*internalpb.RetrieveResults) error) *MockQueryStreamSegmentsServer_Send_Call { - _c.Call.Return(run) - return _c -} - -// SendHeader provides a mock function with given fields: _a0 -func (_m *MockQueryStreamSegmentsServer) SendHeader(_a0 metadata.MD) error { - ret := _m.Called(_a0) - - var r0 error - if rf, ok := ret.Get(0).(func(metadata.MD) error); ok { - r0 = rf(_a0) - } else { - r0 = ret.Error(0) - } - - return r0 -} - -// MockQueryStreamSegmentsServer_SendHeader_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SendHeader' -type MockQueryStreamSegmentsServer_SendHeader_Call struct { - *mock.Call -} - -// SendHeader is a helper method to define mock.On call -// - _a0 metadata.MD -func (_e *MockQueryStreamSegmentsServer_Expecter) SendHeader(_a0 interface{}) *MockQueryStreamSegmentsServer_SendHeader_Call { - return &MockQueryStreamSegmentsServer_SendHeader_Call{Call: _e.mock.On("SendHeader", _a0)} -} - -func (_c *MockQueryStreamSegmentsServer_SendHeader_Call) Run(run func(_a0 metadata.MD)) *MockQueryStreamSegmentsServer_SendHeader_Call { - _c.Call.Run(func(args mock.Arguments) { - run(args[0].(metadata.MD)) - }) - return _c -} - -func (_c *MockQueryStreamSegmentsServer_SendHeader_Call) Return(_a0 error) *MockQueryStreamSegmentsServer_SendHeader_Call { - _c.Call.Return(_a0) - return _c -} - -func (_c *MockQueryStreamSegmentsServer_SendHeader_Call) RunAndReturn(run func(metadata.MD) error) *MockQueryStreamSegmentsServer_SendHeader_Call { - _c.Call.Return(run) - return _c -} - -// SendMsg provides a mock function with given fields: m -func (_m *MockQueryStreamSegmentsServer) SendMsg(m interface{}) error { - ret := _m.Called(m) - - var r0 error - if rf, ok := ret.Get(0).(func(interface{}) error); ok { - r0 = rf(m) - } else { - r0 = ret.Error(0) - } - - return r0 -} - -// MockQueryStreamSegmentsServer_SendMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SendMsg' -type MockQueryStreamSegmentsServer_SendMsg_Call struct { - *mock.Call -} - -// SendMsg is a helper method to define mock.On call -// - m interface{} -func (_e *MockQueryStreamSegmentsServer_Expecter) SendMsg(m interface{}) *MockQueryStreamSegmentsServer_SendMsg_Call { - return &MockQueryStreamSegmentsServer_SendMsg_Call{Call: _e.mock.On("SendMsg", m)} -} - -func (_c *MockQueryStreamSegmentsServer_SendMsg_Call) Run(run func(m interface{})) *MockQueryStreamSegmentsServer_SendMsg_Call { - _c.Call.Run(func(args mock.Arguments) { - run(args[0].(interface{})) - }) - return _c -} - -func (_c *MockQueryStreamSegmentsServer_SendMsg_Call) Return(_a0 error) *MockQueryStreamSegmentsServer_SendMsg_Call { - _c.Call.Return(_a0) - return _c -} - -func (_c *MockQueryStreamSegmentsServer_SendMsg_Call) RunAndReturn(run func(interface{}) error) *MockQueryStreamSegmentsServer_SendMsg_Call { - _c.Call.Return(run) - return _c -} - -// SetHeader provides a mock function with given fields: _a0 -func (_m *MockQueryStreamSegmentsServer) SetHeader(_a0 metadata.MD) error { - ret := _m.Called(_a0) - - var r0 error - if rf, ok := ret.Get(0).(func(metadata.MD) error); ok { - r0 = rf(_a0) - } else { - r0 = ret.Error(0) - } - - return r0 -} - -// MockQueryStreamSegmentsServer_SetHeader_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SetHeader' -type MockQueryStreamSegmentsServer_SetHeader_Call struct { - *mock.Call -} - -// SetHeader is a helper method to define mock.On call -// - _a0 metadata.MD -func (_e *MockQueryStreamSegmentsServer_Expecter) SetHeader(_a0 interface{}) *MockQueryStreamSegmentsServer_SetHeader_Call { - return &MockQueryStreamSegmentsServer_SetHeader_Call{Call: _e.mock.On("SetHeader", _a0)} -} - -func (_c *MockQueryStreamSegmentsServer_SetHeader_Call) Run(run func(_a0 metadata.MD)) *MockQueryStreamSegmentsServer_SetHeader_Call { - _c.Call.Run(func(args mock.Arguments) { - run(args[0].(metadata.MD)) - }) - return _c -} - -func (_c *MockQueryStreamSegmentsServer_SetHeader_Call) Return(_a0 error) *MockQueryStreamSegmentsServer_SetHeader_Call { - _c.Call.Return(_a0) - return _c -} - -func (_c *MockQueryStreamSegmentsServer_SetHeader_Call) RunAndReturn(run func(metadata.MD) error) *MockQueryStreamSegmentsServer_SetHeader_Call { - _c.Call.Return(run) - return _c -} - -// SetTrailer provides a mock function with given fields: _a0 -func (_m *MockQueryStreamSegmentsServer) SetTrailer(_a0 metadata.MD) { - _m.Called(_a0) -} - -// MockQueryStreamSegmentsServer_SetTrailer_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SetTrailer' -type MockQueryStreamSegmentsServer_SetTrailer_Call struct { - *mock.Call -} - -// SetTrailer is a helper method to define mock.On call -// - _a0 metadata.MD -func (_e *MockQueryStreamSegmentsServer_Expecter) SetTrailer(_a0 interface{}) *MockQueryStreamSegmentsServer_SetTrailer_Call { - return &MockQueryStreamSegmentsServer_SetTrailer_Call{Call: _e.mock.On("SetTrailer", _a0)} -} - -func (_c *MockQueryStreamSegmentsServer_SetTrailer_Call) Run(run func(_a0 metadata.MD)) *MockQueryStreamSegmentsServer_SetTrailer_Call { - _c.Call.Run(func(args mock.Arguments) { - run(args[0].(metadata.MD)) - }) - return _c -} - -func (_c *MockQueryStreamSegmentsServer_SetTrailer_Call) Return() *MockQueryStreamSegmentsServer_SetTrailer_Call { - _c.Call.Return() - return _c -} - -func (_c *MockQueryStreamSegmentsServer_SetTrailer_Call) RunAndReturn(run func(metadata.MD)) *MockQueryStreamSegmentsServer_SetTrailer_Call { - _c.Call.Return(run) - return _c -} - -// NewMockQueryStreamSegmentsServer creates a new instance of MockQueryStreamSegmentsServer. 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 NewMockQueryStreamSegmentsServer(t interface { - mock.TestingT - Cleanup(func()) -}) *MockQueryStreamSegmentsServer { - mock := &MockQueryStreamSegmentsServer{} - mock.Mock.Test(t) - - t.Cleanup(func() { mock.AssertExpectations(t) }) - - return mock -} diff --git a/internal/util/streamrpc/mocks/mock_query_stream_server.go b/internal/util/streamrpc/mocks/mock_query_stream_server.go deleted file mode 100644 index bebb58f312..0000000000 --- a/internal/util/streamrpc/mocks/mock_query_stream_server.go +++ /dev/null @@ -1,325 +0,0 @@ -// Code generated by mockery v2.32.4. DO NOT EDIT. - -package mocks - -import ( - context "context" - - internalpb "github.com/milvus-io/milvus/internal/proto/internalpb" - metadata "google.golang.org/grpc/metadata" - - mock "github.com/stretchr/testify/mock" -) - -// MockQueryStreamServer is an autogenerated mock type for the QueryNode_QueryStreamServer type -type MockQueryStreamServer struct { - mock.Mock -} - -type MockQueryStreamServer_Expecter struct { - mock *mock.Mock -} - -func (_m *MockQueryStreamServer) EXPECT() *MockQueryStreamServer_Expecter { - return &MockQueryStreamServer_Expecter{mock: &_m.Mock} -} - -// Context provides a mock function with given fields: -func (_m *MockQueryStreamServer) Context() context.Context { - ret := _m.Called() - - var r0 context.Context - if rf, ok := ret.Get(0).(func() context.Context); ok { - r0 = rf() - } else { - if ret.Get(0) != nil { - r0 = ret.Get(0).(context.Context) - } - } - - return r0 -} - -// MockQueryStreamServer_Context_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Context' -type MockQueryStreamServer_Context_Call struct { - *mock.Call -} - -// Context is a helper method to define mock.On call -func (_e *MockQueryStreamServer_Expecter) Context() *MockQueryStreamServer_Context_Call { - return &MockQueryStreamServer_Context_Call{Call: _e.mock.On("Context")} -} - -func (_c *MockQueryStreamServer_Context_Call) Run(run func()) *MockQueryStreamServer_Context_Call { - _c.Call.Run(func(args mock.Arguments) { - run() - }) - return _c -} - -func (_c *MockQueryStreamServer_Context_Call) Return(_a0 context.Context) *MockQueryStreamServer_Context_Call { - _c.Call.Return(_a0) - return _c -} - -func (_c *MockQueryStreamServer_Context_Call) RunAndReturn(run func() context.Context) *MockQueryStreamServer_Context_Call { - _c.Call.Return(run) - return _c -} - -// RecvMsg provides a mock function with given fields: m -func (_m *MockQueryStreamServer) RecvMsg(m interface{}) error { - ret := _m.Called(m) - - var r0 error - if rf, ok := ret.Get(0).(func(interface{}) error); ok { - r0 = rf(m) - } else { - r0 = ret.Error(0) - } - - return r0 -} - -// MockQueryStreamServer_RecvMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RecvMsg' -type MockQueryStreamServer_RecvMsg_Call struct { - *mock.Call -} - -// RecvMsg is a helper method to define mock.On call -// - m interface{} -func (_e *MockQueryStreamServer_Expecter) RecvMsg(m interface{}) *MockQueryStreamServer_RecvMsg_Call { - return &MockQueryStreamServer_RecvMsg_Call{Call: _e.mock.On("RecvMsg", m)} -} - -func (_c *MockQueryStreamServer_RecvMsg_Call) Run(run func(m interface{})) *MockQueryStreamServer_RecvMsg_Call { - _c.Call.Run(func(args mock.Arguments) { - run(args[0].(interface{})) - }) - return _c -} - -func (_c *MockQueryStreamServer_RecvMsg_Call) Return(_a0 error) *MockQueryStreamServer_RecvMsg_Call { - _c.Call.Return(_a0) - return _c -} - -func (_c *MockQueryStreamServer_RecvMsg_Call) RunAndReturn(run func(interface{}) error) *MockQueryStreamServer_RecvMsg_Call { - _c.Call.Return(run) - return _c -} - -// Send provides a mock function with given fields: _a0 -func (_m *MockQueryStreamServer) Send(_a0 *internalpb.RetrieveResults) error { - ret := _m.Called(_a0) - - var r0 error - if rf, ok := ret.Get(0).(func(*internalpb.RetrieveResults) error); ok { - r0 = rf(_a0) - } else { - r0 = ret.Error(0) - } - - return r0 -} - -// MockQueryStreamServer_Send_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Send' -type MockQueryStreamServer_Send_Call struct { - *mock.Call -} - -// Send is a helper method to define mock.On call -// - _a0 *internalpb.RetrieveResults -func (_e *MockQueryStreamServer_Expecter) Send(_a0 interface{}) *MockQueryStreamServer_Send_Call { - return &MockQueryStreamServer_Send_Call{Call: _e.mock.On("Send", _a0)} -} - -func (_c *MockQueryStreamServer_Send_Call) Run(run func(_a0 *internalpb.RetrieveResults)) *MockQueryStreamServer_Send_Call { - _c.Call.Run(func(args mock.Arguments) { - run(args[0].(*internalpb.RetrieveResults)) - }) - return _c -} - -func (_c *MockQueryStreamServer_Send_Call) Return(_a0 error) *MockQueryStreamServer_Send_Call { - _c.Call.Return(_a0) - return _c -} - -func (_c *MockQueryStreamServer_Send_Call) RunAndReturn(run func(*internalpb.RetrieveResults) error) *MockQueryStreamServer_Send_Call { - _c.Call.Return(run) - return _c -} - -// SendHeader provides a mock function with given fields: _a0 -func (_m *MockQueryStreamServer) SendHeader(_a0 metadata.MD) error { - ret := _m.Called(_a0) - - var r0 error - if rf, ok := ret.Get(0).(func(metadata.MD) error); ok { - r0 = rf(_a0) - } else { - r0 = ret.Error(0) - } - - return r0 -} - -// MockQueryStreamServer_SendHeader_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SendHeader' -type MockQueryStreamServer_SendHeader_Call struct { - *mock.Call -} - -// SendHeader is a helper method to define mock.On call -// - _a0 metadata.MD -func (_e *MockQueryStreamServer_Expecter) SendHeader(_a0 interface{}) *MockQueryStreamServer_SendHeader_Call { - return &MockQueryStreamServer_SendHeader_Call{Call: _e.mock.On("SendHeader", _a0)} -} - -func (_c *MockQueryStreamServer_SendHeader_Call) Run(run func(_a0 metadata.MD)) *MockQueryStreamServer_SendHeader_Call { - _c.Call.Run(func(args mock.Arguments) { - run(args[0].(metadata.MD)) - }) - return _c -} - -func (_c *MockQueryStreamServer_SendHeader_Call) Return(_a0 error) *MockQueryStreamServer_SendHeader_Call { - _c.Call.Return(_a0) - return _c -} - -func (_c *MockQueryStreamServer_SendHeader_Call) RunAndReturn(run func(metadata.MD) error) *MockQueryStreamServer_SendHeader_Call { - _c.Call.Return(run) - return _c -} - -// SendMsg provides a mock function with given fields: m -func (_m *MockQueryStreamServer) SendMsg(m interface{}) error { - ret := _m.Called(m) - - var r0 error - if rf, ok := ret.Get(0).(func(interface{}) error); ok { - r0 = rf(m) - } else { - r0 = ret.Error(0) - } - - return r0 -} - -// MockQueryStreamServer_SendMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SendMsg' -type MockQueryStreamServer_SendMsg_Call struct { - *mock.Call -} - -// SendMsg is a helper method to define mock.On call -// - m interface{} -func (_e *MockQueryStreamServer_Expecter) SendMsg(m interface{}) *MockQueryStreamServer_SendMsg_Call { - return &MockQueryStreamServer_SendMsg_Call{Call: _e.mock.On("SendMsg", m)} -} - -func (_c *MockQueryStreamServer_SendMsg_Call) Run(run func(m interface{})) *MockQueryStreamServer_SendMsg_Call { - _c.Call.Run(func(args mock.Arguments) { - run(args[0].(interface{})) - }) - return _c -} - -func (_c *MockQueryStreamServer_SendMsg_Call) Return(_a0 error) *MockQueryStreamServer_SendMsg_Call { - _c.Call.Return(_a0) - return _c -} - -func (_c *MockQueryStreamServer_SendMsg_Call) RunAndReturn(run func(interface{}) error) *MockQueryStreamServer_SendMsg_Call { - _c.Call.Return(run) - return _c -} - -// SetHeader provides a mock function with given fields: _a0 -func (_m *MockQueryStreamServer) SetHeader(_a0 metadata.MD) error { - ret := _m.Called(_a0) - - var r0 error - if rf, ok := ret.Get(0).(func(metadata.MD) error); ok { - r0 = rf(_a0) - } else { - r0 = ret.Error(0) - } - - return r0 -} - -// MockQueryStreamServer_SetHeader_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SetHeader' -type MockQueryStreamServer_SetHeader_Call struct { - *mock.Call -} - -// SetHeader is a helper method to define mock.On call -// - _a0 metadata.MD -func (_e *MockQueryStreamServer_Expecter) SetHeader(_a0 interface{}) *MockQueryStreamServer_SetHeader_Call { - return &MockQueryStreamServer_SetHeader_Call{Call: _e.mock.On("SetHeader", _a0)} -} - -func (_c *MockQueryStreamServer_SetHeader_Call) Run(run func(_a0 metadata.MD)) *MockQueryStreamServer_SetHeader_Call { - _c.Call.Run(func(args mock.Arguments) { - run(args[0].(metadata.MD)) - }) - return _c -} - -func (_c *MockQueryStreamServer_SetHeader_Call) Return(_a0 error) *MockQueryStreamServer_SetHeader_Call { - _c.Call.Return(_a0) - return _c -} - -func (_c *MockQueryStreamServer_SetHeader_Call) RunAndReturn(run func(metadata.MD) error) *MockQueryStreamServer_SetHeader_Call { - _c.Call.Return(run) - return _c -} - -// SetTrailer provides a mock function with given fields: _a0 -func (_m *MockQueryStreamServer) SetTrailer(_a0 metadata.MD) { - _m.Called(_a0) -} - -// MockQueryStreamServer_SetTrailer_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SetTrailer' -type MockQueryStreamServer_SetTrailer_Call struct { - *mock.Call -} - -// SetTrailer is a helper method to define mock.On call -// - _a0 metadata.MD -func (_e *MockQueryStreamServer_Expecter) SetTrailer(_a0 interface{}) *MockQueryStreamServer_SetTrailer_Call { - return &MockQueryStreamServer_SetTrailer_Call{Call: _e.mock.On("SetTrailer", _a0)} -} - -func (_c *MockQueryStreamServer_SetTrailer_Call) Run(run func(_a0 metadata.MD)) *MockQueryStreamServer_SetTrailer_Call { - _c.Call.Run(func(args mock.Arguments) { - run(args[0].(metadata.MD)) - }) - return _c -} - -func (_c *MockQueryStreamServer_SetTrailer_Call) Return() *MockQueryStreamServer_SetTrailer_Call { - _c.Call.Return() - return _c -} - -func (_c *MockQueryStreamServer_SetTrailer_Call) RunAndReturn(run func(metadata.MD)) *MockQueryStreamServer_SetTrailer_Call { - _c.Call.Return(run) - return _c -} - -// NewMockQueryStreamServer creates a new instance of MockQueryStreamServer. 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 NewMockQueryStreamServer(t interface { - mock.TestingT - Cleanup(func()) -}) *MockQueryStreamServer { - mock := &MockQueryStreamServer{} - mock.Mock.Test(t) - - t.Cleanup(func() { mock.AssertExpectations(t) }) - - return mock -}