milvus/internal/mocks/mock_querynode_client.go

// Code generated by mockery v2.46.0. DO NOT EDIT.

package mocks

import (
	context "context"

	commonpb "github.com/milvus-io/milvus-proto/go-api/v2/commonpb"

	grpc "google.golang.org/grpc"

	internalpb "github.com/milvus-io/milvus/internal/proto/internalpb"

	milvuspb "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"

	mock "github.com/stretchr/testify/mock"

	querypb "github.com/milvus-io/milvus/internal/proto/querypb"
)

// MockQueryNodeClient is an autogenerated mock type for the QueryNodeClient type
type MockQueryNodeClient struct {
	mock.Mock
}

type MockQueryNodeClient_Expecter struct {
	mock *mock.Mock
}

func (_m *MockQueryNodeClient) EXPECT() *MockQueryNodeClient_Expecter {
	return &MockQueryNodeClient_Expecter{mock: &_m.Mock}
}

// Close provides a mock function with given fields:
func (_m *MockQueryNodeClient) Close() error {
	ret := _m.Called()

	if len(ret) == 0 {
		panic("no return value specified for Close")
	}

	var r0 error
	if rf, ok := ret.Get(0).(func() error); ok {
		r0 = rf()
	} else {
		r0 = ret.Error(0)
	}

	return r0
}

// MockQueryNodeClient_Close_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Close'
type MockQueryNodeClient_Close_Call struct {
	*mock.Call
}

// Close is a helper method to define mock.On call
func (_e *MockQueryNodeClient_Expecter) Close() *MockQueryNodeClient_Close_Call {
	return &MockQueryNodeClient_Close_Call{Call: _e.mock.On("Close")}
}

func (_c *MockQueryNodeClient_Close_Call) Run(run func()) *MockQueryNodeClient_Close_Call {
	_c.Call.Run(func(args mock.Arguments) {
		run()
	})
	return _c
}

func (_c *MockQueryNodeClient_Close_Call) Return(_a0 error) *MockQueryNodeClient_Close_Call {
	_c.Call.Return(_a0)
	return _c
}

func (_c *MockQueryNodeClient_Close_Call) RunAndReturn(run func() error) *MockQueryNodeClient_Close_Call {
	_c.Call.Return(run)
	return _c
}

// Delete provides a mock function with given fields: ctx, in, opts
func (_m *MockQueryNodeClient) Delete(ctx context.Context, in *querypb.DeleteRequest, opts ...grpc.CallOption) (*commonpb.Status, error) {
	_va := make([]interface{}, len(opts))
	for _i := range opts {
		_va[_i] = opts[_i]
	}
	var _ca []interface{}
	_ca = append(_ca, ctx, in)
	_ca = append(_ca, _va...)
	ret := _m.Called(_ca...)

	if len(ret) == 0 {
		panic("no return value specified for Delete")
	}

	var r0 *commonpb.Status
	var r1 error
	if rf, ok := ret.Get(0).(func(context.Context, *querypb.DeleteRequest, ...grpc.CallOption) (*commonpb.Status, error)); ok {
		return rf(ctx, in, opts...)
	}
	if rf, ok := ret.Get(0).(func(context.Context, *querypb.DeleteRequest, ...grpc.CallOption) *commonpb.Status); ok {
		r0 = rf(ctx, in, opts...)
	} else {
		if ret.Get(0) != nil {
			r0 = ret.Get(0).(*commonpb.Status)
		}
	}

	if rf, ok := ret.Get(1).(func(context.Context, *querypb.DeleteRequest, ...grpc.CallOption) error); ok {
		r1 = rf(ctx, in, opts...)
	} else {
		r1 = ret.Error(1)
	}

	return r0, r1
}

// MockQueryNodeClient_Delete_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Delete'
type MockQueryNodeClient_Delete_Call struct {
	*mock.Call
}

// Delete is a helper method to define mock.On call
//   - ctx context.Context
//   - in *querypb.DeleteRequest
//   - opts ...grpc.CallOption
func (_e *MockQueryNodeClient_Expecter) Delete(ctx interface{}, in interface{}, opts ...interface{}) *MockQueryNodeClient_Delete_Call {
	return &MockQueryNodeClient_Delete_Call{Call: _e.mock.On("Delete",
		append([]interface{}{ctx, in}, opts...)...)}
}

func (_c *MockQueryNodeClient_Delete_Call) Run(run func(ctx context.Context, in *querypb.DeleteRequest, opts ...grpc.CallOption)) *MockQueryNodeClient_Delete_Call {
	_c.Call.Run(func(args mock.Arguments) {
		variadicArgs := make([]grpc.CallOption, len(args)-2)
		for i, a := range args[2:] {
			if a != nil {
				variadicArgs[i] = a.(grpc.CallOption)
			}
		}
		run(args[0].(context.Context), args[1].(*querypb.DeleteRequest), variadicArgs...)
	})
	return _c
}

func (_c *MockQueryNodeClient_Delete_Call) Return(_a0 *commonpb.Status, _a1 error) *MockQueryNodeClient_Delete_Call {
	_c.Call.Return(_a0, _a1)
	return _c
}

func (_c *MockQueryNodeClient_Delete_Call) RunAndReturn(run func(context.Context, *querypb.DeleteRequest, ...grpc.CallOption) (*commonpb.Status, error)) *MockQueryNodeClient_Delete_Call {
	_c.Call.Return(run)
	return _c
}

// DeleteBatch provides a mock function with given fields: ctx, in, opts
func (_m *MockQueryNodeClient) DeleteBatch(ctx context.Context, in *querypb.DeleteBatchRequest, opts ...grpc.CallOption) (*querypb.DeleteBatchResponse, error) {
	_va := make([]interface{}, len(opts))
	for _i := range opts {
		_va[_i] = opts[_i]
	}
	var _ca []interface{}
	_ca = append(_ca, ctx, in)
	_ca = append(_ca, _va...)
	ret := _m.Called(_ca...)

	if len(ret) == 0 {
		panic("no return value specified for DeleteBatch")
	}

	var r0 *querypb.DeleteBatchResponse
	var r1 error
	if rf, ok := ret.Get(0).(func(context.Context, *querypb.DeleteBatchRequest, ...grpc.CallOption) (*querypb.DeleteBatchResponse, error)); ok {
		return rf(ctx, in, opts...)
	}
	if rf, ok := ret.Get(0).(func(context.Context, *querypb.DeleteBatchRequest, ...grpc.CallOption) *querypb.DeleteBatchResponse); ok {
		r0 = rf(ctx, in, opts...)
	} else {
		if ret.Get(0) != nil {
			r0 = ret.Get(0).(*querypb.DeleteBatchResponse)
		}
	}

	if rf, ok := ret.Get(1).(func(context.Context, *querypb.DeleteBatchRequest, ...grpc.CallOption) error); ok {
		r1 = rf(ctx, in, opts...)
	} else {
		r1 = ret.Error(1)
	}

	return r0, r1
}

// MockQueryNodeClient_DeleteBatch_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'DeleteBatch'
type MockQueryNodeClient_DeleteBatch_Call struct {
	*mock.Call
}

// DeleteBatch is a helper method to define mock.On call
//   - ctx context.Context
//   - in *querypb.DeleteBatchRequest
//   - opts ...grpc.CallOption
func (_e *MockQueryNodeClient_Expecter) DeleteBatch(ctx interface{}, in interface{}, opts ...interface{}) *MockQueryNodeClient_DeleteBatch_Call {
	return &MockQueryNodeClient_DeleteBatch_Call{Call: _e.mock.On("DeleteBatch",
		append([]interface{}{ctx, in}, opts...)...)}
}

func (_c *MockQueryNodeClient_DeleteBatch_Call) Run(run func(ctx context.Context, in *querypb.DeleteBatchRequest, opts ...grpc.CallOption)) *MockQueryNodeClient_DeleteBatch_Call {
	_c.Call.Run(func(args mock.Arguments) {
		variadicArgs := make([]grpc.CallOption, len(args)-2)
		for i, a := range args[2:] {
			if a != nil {
				variadicArgs[i] = a.(grpc.CallOption)
			}
		}
		run(args[0].(context.Context), args[1].(*querypb.DeleteBatchRequest), variadicArgs...)
	})
	return _c
}

func (_c *MockQueryNodeClient_DeleteBatch_Call) Return(_a0 *querypb.DeleteBatchResponse, _a1 error) *MockQueryNodeClient_DeleteBatch_Call {
	_c.Call.Return(_a0, _a1)
	return _c
}

func (_c *MockQueryNodeClient_DeleteBatch_Call) RunAndReturn(run func(context.Context, *querypb.DeleteBatchRequest, ...grpc.CallOption) (*querypb.DeleteBatchResponse, error)) *MockQueryNodeClient_DeleteBatch_Call {
	_c.Call.Return(run)
	return _c
}

// GetComponentStates provides a mock function with given fields: ctx, in, opts
func (_m *MockQueryNodeClient) GetComponentStates(ctx context.Context, in *milvuspb.GetComponentStatesRequest, opts ...grpc.CallOption) (*milvuspb.ComponentStates, error) {
	_va := make([]interface{}, len(opts))
	for _i := range opts {
		_va[_i] = opts[_i]
	}
	var _ca []interface{}
	_ca = append(_ca, ctx, in)
	_ca = append(_ca, _va...)
	ret := _m.Called(_ca...)

	if len(ret) == 0 {
		panic("no return value specified for GetComponentStates")
	}

	var r0 *milvuspb.ComponentStates
	var r1 error
	if rf, ok := ret.Get(0).(func(context.Context, *milvuspb.GetComponentStatesRequest, ...grpc.CallOption) (*milvuspb.ComponentStates, error)); ok {
		return rf(ctx, in, opts...)
	}
	if rf, ok := ret.Get(0).(func(context.Context, *milvuspb.GetComponentStatesRequest, ...grpc.CallOption) *milvuspb.ComponentStates); ok {
		r0 = rf(ctx, in, opts...)
	} else {
		if ret.Get(0) != nil {
			r0 = ret.Get(0).(*milvuspb.ComponentStates)
		}
	}

	if rf, ok := ret.Get(1).(func(context.Context, *milvuspb.GetComponentStatesRequest, ...grpc.CallOption) error); ok {
		r1 = rf(ctx, in, opts...)
	} else {
		r1 = ret.Error(1)
	}

	return r0, r1
}

// MockQueryNodeClient_GetComponentStates_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetComponentStates'
type MockQueryNodeClient_GetComponentStates_Call struct {
	*mock.Call
}

// GetComponentStates is a helper method to define mock.On call
//   - ctx context.Context
//   - in *milvuspb.GetComponentStatesRequest
//   - opts ...grpc.CallOption
func (_e *MockQueryNodeClient_Expecter) GetComponentStates(ctx interface{}, in interface{}, opts ...interface{}) *MockQueryNodeClient_GetComponentStates_Call {
	return &MockQueryNodeClient_GetComponentStates_Call{Call: _e.mock.On("GetComponentStates",
		append([]interface{}{ctx, in}, opts...)...)}
}

func (_c *MockQueryNodeClient_GetComponentStates_Call) Run(run func(ctx context.Context, in *milvuspb.GetComponentStatesRequest, opts ...grpc.CallOption)) *MockQueryNodeClient_GetComponentStates_Call {
	_c.Call.Run(func(args mock.Arguments) {
		variadicArgs := make([]grpc.CallOption, len(args)-2)
		for i, a := range args[2:] {
			if a != nil {
				variadicArgs[i] = a.(grpc.CallOption)
			}
		}
		run(args[0].(context.Context), args[1].(*milvuspb.GetComponentStatesRequest), variadicArgs...)
	})
	return _c
}

func (_c *MockQueryNodeClient_GetComponentStates_Call) Return(_a0 *milvuspb.ComponentStates, _a1 error) *MockQueryNodeClient_GetComponentStates_Call {
	_c.Call.Return(_a0, _a1)
	return _c
}

func (_c *MockQueryNodeClient_GetComponentStates_Call) RunAndReturn(run func(context.Context, *milvuspb.GetComponentStatesRequest, ...grpc.CallOption) (*milvuspb.ComponentStates, error)) *MockQueryNodeClient_GetComponentStates_Call {
	_c.Call.Return(run)
	return _c
}

// GetDataDistribution provides a mock function with given fields: ctx, in, opts
func (_m *MockQueryNodeClient) GetDataDistribution(ctx context.Context, in *querypb.GetDataDistributionRequest, opts ...grpc.CallOption) (*querypb.GetDataDistributionResponse, error) {
	_va := make([]interface{}, len(opts))
	for _i := range opts {
		_va[_i] = opts[_i]
	}
	var _ca []interface{}
	_ca = append(_ca, ctx, in)
	_ca = append(_ca, _va...)
	ret := _m.Called(_ca...)

	if len(ret) == 0 {
		panic("no return value specified for GetDataDistribution")
	}

	var r0 *querypb.GetDataDistributionResponse
	var r1 error
	if rf, ok := ret.Get(0).(func(context.Context, *querypb.GetDataDistributionRequest, ...grpc.CallOption) (*querypb.GetDataDistributionResponse, error)); ok {
		return rf(ctx, in, opts...)
	}
	if rf, ok := ret.Get(0).(func(context.Context, *querypb.GetDataDistributionRequest, ...grpc.CallOption) *querypb.GetDataDistributionResponse); ok {
		r0 = rf(ctx, in, opts...)
	} else {
		if ret.Get(0) != nil {
			r0 = ret.Get(0).(*querypb.GetDataDistributionResponse)
		}
	}

	if rf, ok := ret.Get(1).(func(context.Context, *querypb.GetDataDistributionRequest, ...grpc.CallOption) error); ok {
		r1 = rf(ctx, in, opts...)
	} else {
		r1 = ret.Error(1)
	}

	return r0, r1
}

// MockQueryNodeClient_GetDataDistribution_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetDataDistribution'
type MockQueryNodeClient_GetDataDistribution_Call struct {
	*mock.Call
}

// GetDataDistribution is a helper method to define mock.On call
//   - ctx context.Context
//   - in *querypb.GetDataDistributionRequest
//   - opts ...grpc.CallOption
func (_e *MockQueryNodeClient_Expecter) GetDataDistribution(ctx interface{}, in interface{}, opts ...interface{}) *MockQueryNodeClient_GetDataDistribution_Call {
	return &MockQueryNodeClient_GetDataDistribution_Call{Call: _e.mock.On("GetDataDistribution",
		append([]interface{}{ctx, in}, opts...)...)}
}

func (_c *MockQueryNodeClient_GetDataDistribution_Call) Run(run func(ctx context.Context, in *querypb.GetDataDistributionRequest, opts ...grpc.CallOption)) *MockQueryNodeClient_GetDataDistribution_Call {
	_c.Call.Run(func(args mock.Arguments) {
		variadicArgs := make([]grpc.CallOption, len(args)-2)
		for i, a := range args[2:] {
			if a != nil {
				variadicArgs[i] = a.(grpc.CallOption)
			}
		}
		run(args[0].(context.Context), args[1].(*querypb.GetDataDistributionRequest), variadicArgs...)
	})
	return _c
}

func (_c *MockQueryNodeClient_GetDataDistribution_Call) Return(_a0 *querypb.GetDataDistributionResponse, _a1 error) *MockQueryNodeClient_GetDataDistribution_Call {
	_c.Call.Return(_a0, _a1)
	return _c
}

func (_c *MockQueryNodeClient_GetDataDistribution_Call) RunAndReturn(run func(context.Context, *querypb.GetDataDistributionRequest, ...grpc.CallOption) (*querypb.GetDataDistributionResponse, error)) *MockQueryNodeClient_GetDataDistribution_Call {
	_c.Call.Return(run)
	return _c
}

// GetMetrics provides a mock function with given fields: ctx, in, opts
func (_m *MockQueryNodeClient) GetMetrics(ctx context.Context, in *milvuspb.GetMetricsRequest, opts ...grpc.CallOption) (*milvuspb.GetMetricsResponse, error) {
	_va := make([]interface{}, len(opts))
	for _i := range opts {
		_va[_i] = opts[_i]
	}
	var _ca []interface{}
	_ca = append(_ca, ctx, in)
	_ca = append(_ca, _va...)
	ret := _m.Called(_ca...)

	if len(ret) == 0 {
		panic("no return value specified for GetMetrics")
	}

	var r0 *milvuspb.GetMetricsResponse
	var r1 error
	if rf, ok := ret.Get(0).(func(context.Context, *milvuspb.GetMetricsRequest, ...grpc.CallOption) (*milvuspb.GetMetricsResponse, error)); ok {
		return rf(ctx, in, opts...)
	}
	if rf, ok := ret.Get(0).(func(context.Context, *milvuspb.GetMetricsRequest, ...grpc.CallOption) *milvuspb.GetMetricsResponse); ok {
		r0 = rf(ctx, in, opts...)
	} else {
		if ret.Get(0) != nil {
			r0 = ret.Get(0).(*milvuspb.GetMetricsResponse)
		}
	}

	if rf, ok := ret.Get(1).(func(context.Context, *milvuspb.GetMetricsRequest, ...grpc.CallOption) error); ok {
		r1 = rf(ctx, in, opts...)
	} else {
		r1 = ret.Error(1)
	}

	return r0, r1
}

// MockQueryNodeClient_GetMetrics_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetMetrics'
type MockQueryNodeClient_GetMetrics_Call struct {
	*mock.Call
}

// GetMetrics is a helper method to define mock.On call
//   - ctx context.Context
//   - in *milvuspb.GetMetricsRequest
//   - opts ...grpc.CallOption
func (_e *MockQueryNodeClient_Expecter) GetMetrics(ctx interface{}, in interface{}, opts ...interface{}) *MockQueryNodeClient_GetMetrics_Call {
	return &MockQueryNodeClient_GetMetrics_Call{Call: _e.mock.On("GetMetrics",
		append([]interface{}{ctx, in}, opts...)...)}
}

func (_c *MockQueryNodeClient_GetMetrics_Call) Run(run func(ctx context.Context, in *milvuspb.GetMetricsRequest, opts ...grpc.CallOption)) *MockQueryNodeClient_GetMetrics_Call {
	_c.Call.Run(func(args mock.Arguments) {
		variadicArgs := make([]grpc.CallOption, len(args)-2)
		for i, a := range args[2:] {
			if a != nil {
				variadicArgs[i] = a.(grpc.CallOption)
			}
		}
		run(args[0].(context.Context), args[1].(*milvuspb.GetMetricsRequest), variadicArgs...)
	})
	return _c
}

func (_c *MockQueryNodeClient_GetMetrics_Call) Return(_a0 *milvuspb.GetMetricsResponse, _a1 error) *MockQueryNodeClient_GetMetrics_Call {
	_c.Call.Return(_a0, _a1)
	return _c
}

func (_c *MockQueryNodeClient_GetMetrics_Call) RunAndReturn(run func(context.Context, *milvuspb.GetMetricsRequest, ...grpc.CallOption) (*milvuspb.GetMetricsResponse, error)) *MockQueryNodeClient_GetMetrics_Call {
	_c.Call.Return(run)
	return _c
}

// GetSegmentInfo provides a mock function with given fields: ctx, in, opts
func (_m *MockQueryNodeClient) GetSegmentInfo(ctx context.Context, in *querypb.GetSegmentInfoRequest, opts ...grpc.CallOption) (*querypb.GetSegmentInfoResponse, error) {
	_va := make([]interface{}, len(opts))
	for _i := range opts {
		_va[_i] = opts[_i]
	}
	var _ca []interface{}
	_ca = append(_ca, ctx, in)
	_ca = append(_ca, _va...)
	ret := _m.Called(_ca...)

	if len(ret) == 0 {
		panic("no return value specified for GetSegmentInfo")
	}

	var r0 *querypb.GetSegmentInfoResponse
	var r1 error
	if rf, ok := ret.Get(0).(func(context.Context, *querypb.GetSegmentInfoRequest, ...grpc.CallOption) (*querypb.GetSegmentInfoResponse, error)); ok {
		return rf(ctx, in, opts...)
	}
	if rf, ok := ret.Get(0).(func(context.Context, *querypb.GetSegmentInfoRequest, ...grpc.CallOption) *querypb.GetSegmentInfoResponse); ok {
		r0 = rf(ctx, in, opts...)
	} else {
		if ret.Get(0) != nil {
			r0 = ret.Get(0).(*querypb.GetSegmentInfoResponse)
		}
	}

	if rf, ok := ret.Get(1).(func(context.Context, *querypb.GetSegmentInfoRequest, ...grpc.CallOption) error); ok {
		r1 = rf(ctx, in, opts...)
	} else {
		r1 = ret.Error(1)
	}

	return r0, r1
}

// MockQueryNodeClient_GetSegmentInfo_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetSegmentInfo'
type MockQueryNodeClient_GetSegmentInfo_Call struct {
	*mock.Call
}

// GetSegmentInfo is a helper method to define mock.On call
//   - ctx context.Context
//   - in *querypb.GetSegmentInfoRequest
//   - opts ...grpc.CallOption
func (_e *MockQueryNodeClient_Expecter) GetSegmentInfo(ctx interface{}, in interface{}, opts ...interface{}) *MockQueryNodeClient_GetSegmentInfo_Call {
	return &MockQueryNodeClient_GetSegmentInfo_Call{Call: _e.mock.On("GetSegmentInfo",
		append([]interface{}{ctx, in}, opts...)...)}
}

func (_c *MockQueryNodeClient_GetSegmentInfo_Call) Run(run func(ctx context.Context, in *querypb.GetSegmentInfoRequest, opts ...grpc.CallOption)) *MockQueryNodeClient_GetSegmentInfo_Call {
	_c.Call.Run(func(args mock.Arguments) {
		variadicArgs := make([]grpc.CallOption, len(args)-2)
		for i, a := range args[2:] {
			if a != nil {
				variadicArgs[i] = a.(grpc.CallOption)
			}
		}
		run(args[0].(context.Context), args[1].(*querypb.GetSegmentInfoRequest), variadicArgs...)
	})
	return _c
}

func (_c *MockQueryNodeClient_GetSegmentInfo_Call) Return(_a0 *querypb.GetSegmentInfoResponse, _a1 error) *MockQueryNodeClient_GetSegmentInfo_Call {
	_c.Call.Return(_a0, _a1)
	return _c
}

func (_c *MockQueryNodeClient_GetSegmentInfo_Call) RunAndReturn(run func(context.Context, *querypb.GetSegmentInfoRequest, ...grpc.CallOption) (*querypb.GetSegmentInfoResponse, error)) *MockQueryNodeClient_GetSegmentInfo_Call {
	_c.Call.Return(run)
	return _c
}

// GetStatistics provides a mock function with given fields: ctx, in, opts
func (_m *MockQueryNodeClient) GetStatistics(ctx context.Context, in *querypb.GetStatisticsRequest, opts ...grpc.CallOption) (*internalpb.GetStatisticsResponse, error) {
	_va := make([]interface{}, len(opts))
	for _i := range opts {
		_va[_i] = opts[_i]
	}
	var _ca []interface{}
	_ca = append(_ca, ctx, in)
	_ca = append(_ca, _va...)
	ret := _m.Called(_ca...)

	if len(ret) == 0 {
		panic("no return value specified for GetStatistics")
	}

	var r0 *internalpb.GetStatisticsResponse
	var r1 error
	if rf, ok := ret.Get(0).(func(context.Context, *querypb.GetStatisticsRequest, ...grpc.CallOption) (*internalpb.GetStatisticsResponse, error)); ok {
		return rf(ctx, in, opts...)
	}
	if rf, ok := ret.Get(0).(func(context.Context, *querypb.GetStatisticsRequest, ...grpc.CallOption) *internalpb.GetStatisticsResponse); ok {
		r0 = rf(ctx, in, opts...)
	} else {
		if ret.Get(0) != nil {
			r0 = ret.Get(0).(*internalpb.GetStatisticsResponse)
		}
	}

	if rf, ok := ret.Get(1).(func(context.Context, *querypb.GetStatisticsRequest, ...grpc.CallOption) error); ok {
		r1 = rf(ctx, in, opts...)
	} else {
		r1 = ret.Error(1)
	}

	return r0, r1
}

// MockQueryNodeClient_GetStatistics_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetStatistics'
type MockQueryNodeClient_GetStatistics_Call struct {
	*mock.Call
}

// GetStatistics is a helper method to define mock.On call
//   - ctx context.Context
//   - in *querypb.GetStatisticsRequest
//   - opts ...grpc.CallOption
func (_e *MockQueryNodeClient_Expecter) GetStatistics(ctx interface{}, in interface{}, opts ...interface{}) *MockQueryNodeClient_GetStatistics_Call {
	return &MockQueryNodeClient_GetStatistics_Call{Call: _e.mock.On("GetStatistics",
		append([]interface{}{ctx, in}, opts...)...)}
}

func (_c *MockQueryNodeClient_GetStatistics_Call) Run(run func(ctx context.Context, in *querypb.GetStatisticsRequest, opts ...grpc.CallOption)) *MockQueryNodeClient_GetStatistics_Call {
	_c.Call.Run(func(args mock.Arguments) {
		variadicArgs := make([]grpc.CallOption, len(args)-2)
		for i, a := range args[2:] {
			if a != nil {
				variadicArgs[i] = a.(grpc.CallOption)
			}
		}
		run(args[0].(context.Context), args[1].(*querypb.GetStatisticsRequest), variadicArgs...)
	})
	return _c
}

func (_c *MockQueryNodeClient_GetStatistics_Call) Return(_a0 *internalpb.GetStatisticsResponse, _a1 error) *MockQueryNodeClient_GetStatistics_Call {
	_c.Call.Return(_a0, _a1)
	return _c
}

func (_c *MockQueryNodeClient_GetStatistics_Call) RunAndReturn(run func(context.Context, *querypb.GetStatisticsRequest, ...grpc.CallOption) (*internalpb.GetStatisticsResponse, error)) *MockQueryNodeClient_GetStatistics_Call {
	_c.Call.Return(run)
	return _c
}

// GetStatisticsChannel provides a mock function with given fields: ctx, in, opts
func (_m *MockQueryNodeClient) GetStatisticsChannel(ctx context.Context, in *internalpb.GetStatisticsChannelRequest, opts ...grpc.CallOption) (*milvuspb.StringResponse, error) {
	_va := make([]interface{}, len(opts))
	for _i := range opts {
		_va[_i] = opts[_i]
	}
	var _ca []interface{}
	_ca = append(_ca, ctx, in)
	_ca = append(_ca, _va...)
	ret := _m.Called(_ca...)

	if len(ret) == 0 {
		panic("no return value specified for GetStatisticsChannel")
	}

	var r0 *milvuspb.StringResponse
	var r1 error
	if rf, ok := ret.Get(0).(func(context.Context, *internalpb.GetStatisticsChannelRequest, ...grpc.CallOption) (*milvuspb.StringResponse, error)); ok {
		return rf(ctx, in, opts...)
	}
	if rf, ok := ret.Get(0).(func(context.Context, *internalpb.GetStatisticsChannelRequest, ...grpc.CallOption) *milvuspb.StringResponse); ok {
		r0 = rf(ctx, in, opts...)
	} else {
		if ret.Get(0) != nil {
			r0 = ret.Get(0).(*milvuspb.StringResponse)
		}
	}

	if rf, ok := ret.Get(1).(func(context.Context, *internalpb.GetStatisticsChannelRequest, ...grpc.CallOption) error); ok {
		r1 = rf(ctx, in, opts...)
	} else {
		r1 = ret.Error(1)
	}

	return r0, r1
}

// MockQueryNodeClient_GetStatisticsChannel_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetStatisticsChannel'
type MockQueryNodeClient_GetStatisticsChannel_Call struct {
	*mock.Call
}

// GetStatisticsChannel is a helper method to define mock.On call
//   - ctx context.Context
//   - in *internalpb.GetStatisticsChannelRequest
//   - opts ...grpc.CallOption
func (_e *MockQueryNodeClient_Expecter) GetStatisticsChannel(ctx interface{}, in interface{}, opts ...interface{}) *MockQueryNodeClient_GetStatisticsChannel_Call {
	return &MockQueryNodeClient_GetStatisticsChannel_Call{Call: _e.mock.On("GetStatisticsChannel",
		append([]interface{}{ctx, in}, opts...)...)}
}

func (_c *MockQueryNodeClient_GetStatisticsChannel_Call) Run(run func(ctx context.Context, in *internalpb.GetStatisticsChannelRequest, opts ...grpc.CallOption)) *MockQueryNodeClient_GetStatisticsChannel_Call {
	_c.Call.Run(func(args mock.Arguments) {
		variadicArgs := make([]grpc.CallOption, len(args)-2)
		for i, a := range args[2:] {
			if a != nil {
				variadicArgs[i] = a.(grpc.CallOption)
			}
		}
		run(args[0].(context.Context), args[1].(*internalpb.GetStatisticsChannelRequest), variadicArgs...)
	})
	return _c
}

func (_c *MockQueryNodeClient_GetStatisticsChannel_Call) Return(_a0 *milvuspb.StringResponse, _a1 error) *MockQueryNodeClient_GetStatisticsChannel_Call {
	_c.Call.Return(_a0, _a1)
	return _c
}

func (_c *MockQueryNodeClient_GetStatisticsChannel_Call) RunAndReturn(run func(context.Context, *internalpb.GetStatisticsChannelRequest, ...grpc.CallOption) (*milvuspb.StringResponse, error)) *MockQueryNodeClient_GetStatisticsChannel_Call {
	_c.Call.Return(run)
	return _c
}

// GetTimeTickChannel provides a mock function with given fields: ctx, in, opts
func (_m *MockQueryNodeClient) GetTimeTickChannel(ctx context.Context, in *internalpb.GetTimeTickChannelRequest, opts ...grpc.CallOption) (*milvuspb.StringResponse, error) {
	_va := make([]interface{}, len(opts))
	for _i := range opts {
		_va[_i] = opts[_i]
	}
	var _ca []interface{}
	_ca = append(_ca, ctx, in)
	_ca = append(_ca, _va...)
	ret := _m.Called(_ca...)

	if len(ret) == 0 {
		panic("no return value specified for GetTimeTickChannel")
	}

	var r0 *milvuspb.StringResponse
	var r1 error
	if rf, ok := ret.Get(0).(func(context.Context, *internalpb.GetTimeTickChannelRequest, ...grpc.CallOption) (*milvuspb.StringResponse, error)); ok {
		return rf(ctx, in, opts...)
	}
	if rf, ok := ret.Get(0).(func(context.Context, *internalpb.GetTimeTickChannelRequest, ...grpc.CallOption) *milvuspb.StringResponse); ok {
		r0 = rf(ctx, in, opts...)
	} else {
		if ret.Get(0) != nil {
			r0 = ret.Get(0).(*milvuspb.StringResponse)
		}
	}

	if rf, ok := ret.Get(1).(func(context.Context, *internalpb.GetTimeTickChannelRequest, ...grpc.CallOption) error); ok {
		r1 = rf(ctx, in, opts...)
	} else {
		r1 = ret.Error(1)
	}

	return r0, r1
}

// MockQueryNodeClient_GetTimeTickChannel_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetTimeTickChannel'
type MockQueryNodeClient_GetTimeTickChannel_Call struct {
	*mock.Call
}

// GetTimeTickChannel is a helper method to define mock.On call
//   - ctx context.Context
//   - in *internalpb.GetTimeTickChannelRequest
//   - opts ...grpc.CallOption
func (_e *MockQueryNodeClient_Expecter) GetTimeTickChannel(ctx interface{}, in interface{}, opts ...interface{}) *MockQueryNodeClient_GetTimeTickChannel_Call {
	return &MockQueryNodeClient_GetTimeTickChannel_Call{Call: _e.mock.On("GetTimeTickChannel",
		append([]interface{}{ctx, in}, opts...)...)}
}

func (_c *MockQueryNodeClient_GetTimeTickChannel_Call) Run(run func(ctx context.Context, in *internalpb.GetTimeTickChannelRequest, opts ...grpc.CallOption)) *MockQueryNodeClient_GetTimeTickChannel_Call {
	_c.Call.Run(func(args mock.Arguments) {
		variadicArgs := make([]grpc.CallOption, len(args)-2)
		for i, a := range args[2:] {
			if a != nil {
				variadicArgs[i] = a.(grpc.CallOption)
			}
		}
		run(args[0].(context.Context), args[1].(*internalpb.GetTimeTickChannelRequest), variadicArgs...)
	})
	return _c
}

func (_c *MockQueryNodeClient_GetTimeTickChannel_Call) Return(_a0 *milvuspb.StringResponse, _a1 error) *MockQueryNodeClient_GetTimeTickChannel_Call {
	_c.Call.Return(_a0, _a1)
	return _c
}

func (_c *MockQueryNodeClient_GetTimeTickChannel_Call) RunAndReturn(run func(context.Context, *internalpb.GetTimeTickChannelRequest, ...grpc.CallOption) (*milvuspb.StringResponse, error)) *MockQueryNodeClient_GetTimeTickChannel_Call {
	_c.Call.Return(run)
	return _c
}

// LoadPartitions provides a mock function with given fields: ctx, in, opts
func (_m *MockQueryNodeClient) LoadPartitions(ctx context.Context, in *querypb.LoadPartitionsRequest, opts ...grpc.CallOption) (*commonpb.Status, error) {
	_va := make([]interface{}, len(opts))
	for _i := range opts {
		_va[_i] = opts[_i]
	}
	var _ca []interface{}
	_ca = append(_ca, ctx, in)
	_ca = append(_ca, _va...)
	ret := _m.Called(_ca...)

	if len(ret) == 0 {
		panic("no return value specified for LoadPartitions")
	}

	var r0 *commonpb.Status
	var r1 error
	if rf, ok := ret.Get(0).(func(context.Context, *querypb.LoadPartitionsRequest, ...grpc.CallOption) (*commonpb.Status, error)); ok {
		return rf(ctx, in, opts...)
	}
	if rf, ok := ret.Get(0).(func(context.Context, *querypb.LoadPartitionsRequest, ...grpc.CallOption) *commonpb.Status); ok {
		r0 = rf(ctx, in, opts...)
	} else {
		if ret.Get(0) != nil {
			r0 = ret.Get(0).(*commonpb.Status)
		}
	}

	if rf, ok := ret.Get(1).(func(context.Context, *querypb.LoadPartitionsRequest, ...grpc.CallOption) error); ok {
		r1 = rf(ctx, in, opts...)
	} else {
		r1 = ret.Error(1)
	}

	return r0, r1
}

// MockQueryNodeClient_LoadPartitions_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'LoadPartitions'
type MockQueryNodeClient_LoadPartitions_Call struct {
	*mock.Call
}

// LoadPartitions is a helper method to define mock.On call
//   - ctx context.Context
//   - in *querypb.LoadPartitionsRequest
//   - opts ...grpc.CallOption
func (_e *MockQueryNodeClient_Expecter) LoadPartitions(ctx interface{}, in interface{}, opts ...interface{}) *MockQueryNodeClient_LoadPartitions_Call {
	return &MockQueryNodeClient_LoadPartitions_Call{Call: _e.mock.On("LoadPartitions",
		append([]interface{}{ctx, in}, opts...)...)}
}

func (_c *MockQueryNodeClient_LoadPartitions_Call) Run(run func(ctx context.Context, in *querypb.LoadPartitionsRequest, opts ...grpc.CallOption)) *MockQueryNodeClient_LoadPartitions_Call {
	_c.Call.Run(func(args mock.Arguments) {
		variadicArgs := make([]grpc.CallOption, len(args)-2)
		for i, a := range args[2:] {
			if a != nil {
				variadicArgs[i] = a.(grpc.CallOption)
			}
		}
		run(args[0].(context.Context), args[1].(*querypb.LoadPartitionsRequest), variadicArgs...)
	})
	return _c
}

func (_c *MockQueryNodeClient_LoadPartitions_Call) Return(_a0 *commonpb.Status, _a1 error) *MockQueryNodeClient_LoadPartitions_Call {
	_c.Call.Return(_a0, _a1)
	return _c
}

func (_c *MockQueryNodeClient_LoadPartitions_Call) RunAndReturn(run func(context.Context, *querypb.LoadPartitionsRequest, ...grpc.CallOption) (*commonpb.Status, error)) *MockQueryNodeClient_LoadPartitions_Call {
	_c.Call.Return(run)
	return _c
}

// LoadSegments provides a mock function with given fields: ctx, in, opts
func (_m *MockQueryNodeClient) LoadSegments(ctx context.Context, in *querypb.LoadSegmentsRequest, opts ...grpc.CallOption) (*commonpb.Status, error) {
	_va := make([]interface{}, len(opts))
	for _i := range opts {
		_va[_i] = opts[_i]
	}
	var _ca []interface{}
	_ca = append(_ca, ctx, in)
	_ca = append(_ca, _va...)
	ret := _m.Called(_ca...)

	if len(ret) == 0 {
		panic("no return value specified for LoadSegments")
	}

	var r0 *commonpb.Status
	var r1 error
	if rf, ok := ret.Get(0).(func(context.Context, *querypb.LoadSegmentsRequest, ...grpc.CallOption) (*commonpb.Status, error)); ok {
		return rf(ctx, in, opts...)
	}
	if rf, ok := ret.Get(0).(func(context.Context, *querypb.LoadSegmentsRequest, ...grpc.CallOption) *commonpb.Status); ok {
		r0 = rf(ctx, in, opts...)
	} else {
		if ret.Get(0) != nil {
			r0 = ret.Get(0).(*commonpb.Status)
		}
	}

	if rf, ok := ret.Get(1).(func(context.Context, *querypb.LoadSegmentsRequest, ...grpc.CallOption) error); ok {
		r1 = rf(ctx, in, opts...)
	} else {
		r1 = ret.Error(1)
	}

	return r0, r1
}

// MockQueryNodeClient_LoadSegments_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'LoadSegments'
type MockQueryNodeClient_LoadSegments_Call struct {
	*mock.Call
}

// LoadSegments is a helper method to define mock.On call
//   - ctx context.Context
//   - in *querypb.LoadSegmentsRequest
//   - opts ...grpc.CallOption
func (_e *MockQueryNodeClient_Expecter) LoadSegments(ctx interface{}, in interface{}, opts ...interface{}) *MockQueryNodeClient_LoadSegments_Call {
	return &MockQueryNodeClient_LoadSegments_Call{Call: _e.mock.On("LoadSegments",
		append([]interface{}{ctx, in}, opts...)...)}
}

func (_c *MockQueryNodeClient_LoadSegments_Call) Run(run func(ctx context.Context, in *querypb.LoadSegmentsRequest, opts ...grpc.CallOption)) *MockQueryNodeClient_LoadSegments_Call {
	_c.Call.Run(func(args mock.Arguments) {
		variadicArgs := make([]grpc.CallOption, len(args)-2)
		for i, a := range args[2:] {
			if a != nil {
				variadicArgs[i] = a.(grpc.CallOption)
			}
		}
		run(args[0].(context.Context), args[1].(*querypb.LoadSegmentsRequest), variadicArgs...)
	})
	return _c
}

func (_c *MockQueryNodeClient_LoadSegments_Call) Return(_a0 *commonpb.Status, _a1 error) *MockQueryNodeClient_LoadSegments_Call {
	_c.Call.Return(_a0, _a1)
	return _c
}

func (_c *MockQueryNodeClient_LoadSegments_Call) RunAndReturn(run func(context.Context, *querypb.LoadSegmentsRequest, ...grpc.CallOption) (*commonpb.Status, error)) *MockQueryNodeClient_LoadSegments_Call {
	_c.Call.Return(run)
	return _c
}

// Query provides a mock function with given fields: ctx, in, opts
func (_m *MockQueryNodeClient) Query(ctx context.Context, in *querypb.QueryRequest, opts ...grpc.CallOption) (*internalpb.RetrieveResults, error) {
	_va := make([]interface{}, len(opts))
	for _i := range opts {
		_va[_i] = opts[_i]
	}
	var _ca []interface{}
	_ca = append(_ca, ctx, in)
	_ca = append(_ca, _va...)
	ret := _m.Called(_ca...)

	if len(ret) == 0 {
		panic("no return value specified for Query")
	}

	var r0 *internalpb.RetrieveResults
	var r1 error
	if rf, ok := ret.Get(0).(func(context.Context, *querypb.QueryRequest, ...grpc.CallOption) (*internalpb.RetrieveResults, error)); ok {
		return rf(ctx, in, opts...)
	}
	if rf, ok := ret.Get(0).(func(context.Context, *querypb.QueryRequest, ...grpc.CallOption) *internalpb.RetrieveResults); ok {
		r0 = rf(ctx, in, opts...)
	} else {
		if ret.Get(0) != nil {
			r0 = ret.Get(0).(*internalpb.RetrieveResults)
		}
	}

	if rf, ok := ret.Get(1).(func(context.Context, *querypb.QueryRequest, ...grpc.CallOption) error); ok {
		r1 = rf(ctx, in, opts...)
	} else {
		r1 = ret.Error(1)
	}

	return r0, r1
}

// MockQueryNodeClient_Query_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Query'
type MockQueryNodeClient_Query_Call struct {
	*mock.Call
}

// Query is a helper method to define mock.On call
//   - ctx context.Context
//   - in *querypb.QueryRequest
//   - opts ...grpc.CallOption
func (_e *MockQueryNodeClient_Expecter) Query(ctx interface{}, in interface{}, opts ...interface{}) *MockQueryNodeClient_Query_Call {
	return &MockQueryNodeClient_Query_Call{Call: _e.mock.On("Query",
		append([]interface{}{ctx, in}, opts...)...)}
}

func (_c *MockQueryNodeClient_Query_Call) Run(run func(ctx context.Context, in *querypb.QueryRequest, opts ...grpc.CallOption)) *MockQueryNodeClient_Query_Call {
	_c.Call.Run(func(args mock.Arguments) {
		variadicArgs := make([]grpc.CallOption, len(args)-2)
		for i, a := range args[2:] {
			if a != nil {
				variadicArgs[i] = a.(grpc.CallOption)
			}
		}
		run(args[0].(context.Context), args[1].(*querypb.QueryRequest), variadicArgs...)
	})
	return _c
}

func (_c *MockQueryNodeClient_Query_Call) Return(_a0 *internalpb.RetrieveResults, _a1 error) *MockQueryNodeClient_Query_Call {
	_c.Call.Return(_a0, _a1)
	return _c
}

func (_c *MockQueryNodeClient_Query_Call) RunAndReturn(run func(context.Context, *querypb.QueryRequest, ...grpc.CallOption) (*internalpb.RetrieveResults, error)) *MockQueryNodeClient_Query_Call {
	_c.Call.Return(run)
	return _c
}

// QuerySegments provides a mock function with given fields: ctx, in, opts
func (_m *MockQueryNodeClient) QuerySegments(ctx context.Context, in *querypb.QueryRequest, opts ...grpc.CallOption) (*internalpb.RetrieveResults, error) {
	_va := make([]interface{}, len(opts))
	for _i := range opts {
		_va[_i] = opts[_i]
	}
	var _ca []interface{}
	_ca = append(_ca, ctx, in)
	_ca = append(_ca, _va...)
	ret := _m.Called(_ca...)

	if len(ret) == 0 {
		panic("no return value specified for QuerySegments")
	}

	var r0 *internalpb.RetrieveResults
	var r1 error
	if rf, ok := ret.Get(0).(func(context.Context, *querypb.QueryRequest, ...grpc.CallOption) (*internalpb.RetrieveResults, error)); ok {
		return rf(ctx, in, opts...)
	}
	if rf, ok := ret.Get(0).(func(context.Context, *querypb.QueryRequest, ...grpc.CallOption) *internalpb.RetrieveResults); ok {
		r0 = rf(ctx, in, opts...)
	} else {
		if ret.Get(0) != nil {
			r0 = ret.Get(0).(*internalpb.RetrieveResults)
		}
	}

	if rf, ok := ret.Get(1).(func(context.Context, *querypb.QueryRequest, ...grpc.CallOption) error); ok {
		r1 = rf(ctx, in, opts...)
	} else {
		r1 = ret.Error(1)
	}

	return r0, r1
}

// MockQueryNodeClient_QuerySegments_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'QuerySegments'
type MockQueryNodeClient_QuerySegments_Call struct {
	*mock.Call
}

// QuerySegments is a helper method to define mock.On call
//   - ctx context.Context
//   - in *querypb.QueryRequest
//   - opts ...grpc.CallOption
func (_e *MockQueryNodeClient_Expecter) QuerySegments(ctx interface{}, in interface{}, opts ...interface{}) *MockQueryNodeClient_QuerySegments_Call {
	return &MockQueryNodeClient_QuerySegments_Call{Call: _e.mock.On("QuerySegments",
		append([]interface{}{ctx, in}, opts...)...)}
}

func (_c *MockQueryNodeClient_QuerySegments_Call) Run(run func(ctx context.Context, in *querypb.QueryRequest, opts ...grpc.CallOption)) *MockQueryNodeClient_QuerySegments_Call {
	_c.Call.Run(func(args mock.Arguments) {
		variadicArgs := make([]grpc.CallOption, len(args)-2)
		for i, a := range args[2:] {
			if a != nil {
				variadicArgs[i] = a.(grpc.CallOption)
			}
		}
		run(args[0].(context.Context), args[1].(*querypb.QueryRequest), variadicArgs...)
	})
	return _c
}

func (_c *MockQueryNodeClient_QuerySegments_Call) Return(_a0 *internalpb.RetrieveResults, _a1 error) *MockQueryNodeClient_QuerySegments_Call {
	_c.Call.Return(_a0, _a1)
	return _c
}

func (_c *MockQueryNodeClient_QuerySegments_Call) RunAndReturn(run func(context.Context, *querypb.QueryRequest, ...grpc.CallOption) (*internalpb.RetrieveResults, error)) *MockQueryNodeClient_QuerySegments_Call {
	_c.Call.Return(run)
	return _c
}

// QueryStream provides a mock function with given fields: ctx, in, opts
func (_m *MockQueryNodeClient) QueryStream(ctx context.Context, in *querypb.QueryRequest, opts ...grpc.CallOption) (querypb.QueryNode_QueryStreamClient, error) {
	_va := make([]interface{}, len(opts))
	for _i := range opts {
		_va[_i] = opts[_i]
	}
	var _ca []interface{}
	_ca = append(_ca, ctx, in)
	_ca = append(_ca, _va...)
	ret := _m.Called(_ca...)

	if len(ret) == 0 {
		panic("no return value specified for QueryStream")
	}

	var r0 querypb.QueryNode_QueryStreamClient
	var r1 error
	if rf, ok := ret.Get(0).(func(context.Context, *querypb.QueryRequest, ...grpc.CallOption) (querypb.QueryNode_QueryStreamClient, error)); ok {
		return rf(ctx, in, opts...)
	}
	if rf, ok := ret.Get(0).(func(context.Context, *querypb.QueryRequest, ...grpc.CallOption) querypb.QueryNode_QueryStreamClient); ok {
		r0 = rf(ctx, in, opts...)
	} else {
		if ret.Get(0) != nil {
			r0 = ret.Get(0).(querypb.QueryNode_QueryStreamClient)
		}
	}

	if rf, ok := ret.Get(1).(func(context.Context, *querypb.QueryRequest, ...grpc.CallOption) error); ok {
		r1 = rf(ctx, in, opts...)
	} else {
		r1 = ret.Error(1)
	}

	return r0, r1
}

// MockQueryNodeClient_QueryStream_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'QueryStream'
type MockQueryNodeClient_QueryStream_Call struct {
	*mock.Call
}

// QueryStream is a helper method to define mock.On call
//   - ctx context.Context
//   - in *querypb.QueryRequest
//   - opts ...grpc.CallOption
func (_e *MockQueryNodeClient_Expecter) QueryStream(ctx interface{}, in interface{}, opts ...interface{}) *MockQueryNodeClient_QueryStream_Call {
	return &MockQueryNodeClient_QueryStream_Call{Call: _e.mock.On("QueryStream",
		append([]interface{}{ctx, in}, opts...)...)}
}

func (_c *MockQueryNodeClient_QueryStream_Call) Run(run func(ctx context.Context, in *querypb.QueryRequest, opts ...grpc.CallOption)) *MockQueryNodeClient_QueryStream_Call {
	_c.Call.Run(func(args mock.Arguments) {
		variadicArgs := make([]grpc.CallOption, len(args)-2)
		for i, a := range args[2:] {
			if a != nil {
				variadicArgs[i] = a.(grpc.CallOption)
			}
		}
		run(args[0].(context.Context), args[1].(*querypb.QueryRequest), variadicArgs...)
	})
	return _c
}

func (_c *MockQueryNodeClient_QueryStream_Call) Return(_a0 querypb.QueryNode_QueryStreamClient, _a1 error) *MockQueryNodeClient_QueryStream_Call {
	_c.Call.Return(_a0, _a1)
	return _c
}

func (_c *MockQueryNodeClient_QueryStream_Call) RunAndReturn(run func(context.Context, *querypb.QueryRequest, ...grpc.CallOption) (querypb.QueryNode_QueryStreamClient, error)) *MockQueryNodeClient_QueryStream_Call {
	_c.Call.Return(run)
	return _c
}

// QueryStreamSegments provides a mock function with given fields: ctx, in, opts
func (_m *MockQueryNodeClient) QueryStreamSegments(ctx context.Context, in *querypb.QueryRequest, opts ...grpc.CallOption) (querypb.QueryNode_QueryStreamSegmentsClient, error) {
	_va := make([]interface{}, len(opts))
	for _i := range opts {
		_va[_i] = opts[_i]
	}
	var _ca []interface{}
	_ca = append(_ca, ctx, in)
	_ca = append(_ca, _va...)
	ret := _m.Called(_ca...)

	if len(ret) == 0 {
		panic("no return value specified for QueryStreamSegments")
	}

	var r0 querypb.QueryNode_QueryStreamSegmentsClient
	var r1 error
	if rf, ok := ret.Get(0).(func(context.Context, *querypb.QueryRequest, ...grpc.CallOption) (querypb.QueryNode_QueryStreamSegmentsClient, error)); ok {
		return rf(ctx, in, opts...)
	}
	if rf, ok := ret.Get(0).(func(context.Context, *querypb.QueryRequest, ...grpc.CallOption) querypb.QueryNode_QueryStreamSegmentsClient); ok {
		r0 = rf(ctx, in, opts...)
	} else {
		if ret.Get(0) != nil {
			r0 = ret.Get(0).(querypb.QueryNode_QueryStreamSegmentsClient)
		}
	}

	if rf, ok := ret.Get(1).(func(context.Context, *querypb.QueryRequest, ...grpc.CallOption) error); ok {
		r1 = rf(ctx, in, opts...)
	} else {
		r1 = ret.Error(1)
	}

	return r0, r1
}

// MockQueryNodeClient_QueryStreamSegments_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'QueryStreamSegments'
type MockQueryNodeClient_QueryStreamSegments_Call struct {
	*mock.Call
}

// QueryStreamSegments is a helper method to define mock.On call
//   - ctx context.Context
//   - in *querypb.QueryRequest
//   - opts ...grpc.CallOption
func (_e *MockQueryNodeClient_Expecter) QueryStreamSegments(ctx interface{}, in interface{}, opts ...interface{}) *MockQueryNodeClient_QueryStreamSegments_Call {
	return &MockQueryNodeClient_QueryStreamSegments_Call{Call: _e.mock.On("QueryStreamSegments",
		append([]interface{}{ctx, in}, opts...)...)}
}

func (_c *MockQueryNodeClient_QueryStreamSegments_Call) Run(run func(ctx context.Context, in *querypb.QueryRequest, opts ...grpc.CallOption)) *MockQueryNodeClient_QueryStreamSegments_Call {
	_c.Call.Run(func(args mock.Arguments) {
		variadicArgs := make([]grpc.CallOption, len(args)-2)
		for i, a := range args[2:] {
			if a != nil {
				variadicArgs[i] = a.(grpc.CallOption)
			}
		}
		run(args[0].(context.Context), args[1].(*querypb.QueryRequest), variadicArgs...)
	})
	return _c
}

func (_c *MockQueryNodeClient_QueryStreamSegments_Call) Return(_a0 querypb.QueryNode_QueryStreamSegmentsClient, _a1 error) *MockQueryNodeClient_QueryStreamSegments_Call {
	_c.Call.Return(_a0, _a1)
	return _c
}

func (_c *MockQueryNodeClient_QueryStreamSegments_Call) RunAndReturn(run func(context.Context, *querypb.QueryRequest, ...grpc.CallOption) (querypb.QueryNode_QueryStreamSegmentsClient, error)) *MockQueryNodeClient_QueryStreamSegments_Call {
	_c.Call.Return(run)
	return _c
}

// ReleaseCollection provides a mock function with given fields: ctx, in, opts
func (_m *MockQueryNodeClient) ReleaseCollection(ctx context.Context, in *querypb.ReleaseCollectionRequest, opts ...grpc.CallOption) (*commonpb.Status, error) {
	_va := make([]interface{}, len(opts))
	for _i := range opts {
		_va[_i] = opts[_i]
	}
	var _ca []interface{}
	_ca = append(_ca, ctx, in)
	_ca = append(_ca, _va...)
	ret := _m.Called(_ca...)

	if len(ret) == 0 {
		panic("no return value specified for ReleaseCollection")
	}

	var r0 *commonpb.Status
	var r1 error
	if rf, ok := ret.Get(0).(func(context.Context, *querypb.ReleaseCollectionRequest, ...grpc.CallOption) (*commonpb.Status, error)); ok {
		return rf(ctx, in, opts...)
	}
	if rf, ok := ret.Get(0).(func(context.Context, *querypb.ReleaseCollectionRequest, ...grpc.CallOption) *commonpb.Status); ok {
		r0 = rf(ctx, in, opts...)
	} else {
		if ret.Get(0) != nil {
			r0 = ret.Get(0).(*commonpb.Status)
		}
	}

	if rf, ok := ret.Get(1).(func(context.Context, *querypb.ReleaseCollectionRequest, ...grpc.CallOption) error); ok {
		r1 = rf(ctx, in, opts...)
	} else {
		r1 = ret.Error(1)
	}

	return r0, r1
}

// MockQueryNodeClient_ReleaseCollection_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'ReleaseCollection'
type MockQueryNodeClient_ReleaseCollection_Call struct {
	*mock.Call
}

// ReleaseCollection is a helper method to define mock.On call
//   - ctx context.Context
//   - in *querypb.ReleaseCollectionRequest
//   - opts ...grpc.CallOption
func (_e *MockQueryNodeClient_Expecter) ReleaseCollection(ctx interface{}, in interface{}, opts ...interface{}) *MockQueryNodeClient_ReleaseCollection_Call {
	return &MockQueryNodeClient_ReleaseCollection_Call{Call: _e.mock.On("ReleaseCollection",
		append([]interface{}{ctx, in}, opts...)...)}
}

func (_c *MockQueryNodeClient_ReleaseCollection_Call) Run(run func(ctx context.Context, in *querypb.ReleaseCollectionRequest, opts ...grpc.CallOption)) *MockQueryNodeClient_ReleaseCollection_Call {
	_c.Call.Run(func(args mock.Arguments) {
		variadicArgs := make([]grpc.CallOption, len(args)-2)
		for i, a := range args[2:] {
			if a != nil {
				variadicArgs[i] = a.(grpc.CallOption)
			}
		}
		run(args[0].(context.Context), args[1].(*querypb.ReleaseCollectionRequest), variadicArgs...)
	})
	return _c
}

func (_c *MockQueryNodeClient_ReleaseCollection_Call) Return(_a0 *commonpb.Status, _a1 error) *MockQueryNodeClient_ReleaseCollection_Call {
	_c.Call.Return(_a0, _a1)
	return _c
}

func (_c *MockQueryNodeClient_ReleaseCollection_Call) RunAndReturn(run func(context.Context, *querypb.ReleaseCollectionRequest, ...grpc.CallOption) (*commonpb.Status, error)) *MockQueryNodeClient_ReleaseCollection_Call {
	_c.Call.Return(run)
	return _c
}

// ReleasePartitions provides a mock function with given fields: ctx, in, opts
func (_m *MockQueryNodeClient) ReleasePartitions(ctx context.Context, in *querypb.ReleasePartitionsRequest, opts ...grpc.CallOption) (*commonpb.Status, error) {
	_va := make([]interface{}, len(opts))
	for _i := range opts {
		_va[_i] = opts[_i]
	}
	var _ca []interface{}
	_ca = append(_ca, ctx, in)
	_ca = append(_ca, _va...)
	ret := _m.Called(_ca...)

	if len(ret) == 0 {
		panic("no return value specified for ReleasePartitions")
	}

	var r0 *commonpb.Status
	var r1 error
	if rf, ok := ret.Get(0).(func(context.Context, *querypb.ReleasePartitionsRequest, ...grpc.CallOption) (*commonpb.Status, error)); ok {
		return rf(ctx, in, opts...)
	}
	if rf, ok := ret.Get(0).(func(context.Context, *querypb.ReleasePartitionsRequest, ...grpc.CallOption) *commonpb.Status); ok {
		r0 = rf(ctx, in, opts...)
	} else {
		if ret.Get(0) != nil {
			r0 = ret.Get(0).(*commonpb.Status)
		}
	}

	if rf, ok := ret.Get(1).(func(context.Context, *querypb.ReleasePartitionsRequest, ...grpc.CallOption) error); ok {
		r1 = rf(ctx, in, opts...)
	} else {
		r1 = ret.Error(1)
	}

	return r0, r1
}

// MockQueryNodeClient_ReleasePartitions_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'ReleasePartitions'
type MockQueryNodeClient_ReleasePartitions_Call struct {
	*mock.Call
}

// ReleasePartitions is a helper method to define mock.On call
//   - ctx context.Context
//   - in *querypb.ReleasePartitionsRequest
//   - opts ...grpc.CallOption
func (_e *MockQueryNodeClient_Expecter) ReleasePartitions(ctx interface{}, in interface{}, opts ...interface{}) *MockQueryNodeClient_ReleasePartitions_Call {
	return &MockQueryNodeClient_ReleasePartitions_Call{Call: _e.mock.On("ReleasePartitions",
		append([]interface{}{ctx, in}, opts...)...)}
}

func (_c *MockQueryNodeClient_ReleasePartitions_Call) Run(run func(ctx context.Context, in *querypb.ReleasePartitionsRequest, opts ...grpc.CallOption)) *MockQueryNodeClient_ReleasePartitions_Call {
	_c.Call.Run(func(args mock.Arguments) {
		variadicArgs := make([]grpc.CallOption, len(args)-2)
		for i, a := range args[2:] {
			if a != nil {
				variadicArgs[i] = a.(grpc.CallOption)
			}
		}
		run(args[0].(context.Context), args[1].(*querypb.ReleasePartitionsRequest), variadicArgs...)
	})
	return _c
}

func (_c *MockQueryNodeClient_ReleasePartitions_Call) Return(_a0 *commonpb.Status, _a1 error) *MockQueryNodeClient_ReleasePartitions_Call {
	_c.Call.Return(_a0, _a1)
	return _c
}

func (_c *MockQueryNodeClient_ReleasePartitions_Call) RunAndReturn(run func(context.Context, *querypb.ReleasePartitionsRequest, ...grpc.CallOption) (*commonpb.Status, error)) *MockQueryNodeClient_ReleasePartitions_Call {
	_c.Call.Return(run)
	return _c
}

// ReleaseSegments provides a mock function with given fields: ctx, in, opts
func (_m *MockQueryNodeClient) ReleaseSegments(ctx context.Context, in *querypb.ReleaseSegmentsRequest, opts ...grpc.CallOption) (*commonpb.Status, error) {
	_va := make([]interface{}, len(opts))
	for _i := range opts {
		_va[_i] = opts[_i]
	}
	var _ca []interface{}
	_ca = append(_ca, ctx, in)
	_ca = append(_ca, _va...)
	ret := _m.Called(_ca...)

	if len(ret) == 0 {
		panic("no return value specified for ReleaseSegments")
	}

	var r0 *commonpb.Status
	var r1 error
	if rf, ok := ret.Get(0).(func(context.Context, *querypb.ReleaseSegmentsRequest, ...grpc.CallOption) (*commonpb.Status, error)); ok {
		return rf(ctx, in, opts...)
	}
	if rf, ok := ret.Get(0).(func(context.Context, *querypb.ReleaseSegmentsRequest, ...grpc.CallOption) *commonpb.Status); ok {
		r0 = rf(ctx, in, opts...)
	} else {
		if ret.Get(0) != nil {
			r0 = ret.Get(0).(*commonpb.Status)
		}
	}

	if rf, ok := ret.Get(1).(func(context.Context, *querypb.ReleaseSegmentsRequest, ...grpc.CallOption) error); ok {
		r1 = rf(ctx, in, opts...)
	} else {
		r1 = ret.Error(1)
	}

	return r0, r1
}

// MockQueryNodeClient_ReleaseSegments_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'ReleaseSegments'
type MockQueryNodeClient_ReleaseSegments_Call struct {
	*mock.Call
}

// ReleaseSegments is a helper method to define mock.On call
//   - ctx context.Context
//   - in *querypb.ReleaseSegmentsRequest
//   - opts ...grpc.CallOption
func (_e *MockQueryNodeClient_Expecter) ReleaseSegments(ctx interface{}, in interface{}, opts ...interface{}) *MockQueryNodeClient_ReleaseSegments_Call {
	return &MockQueryNodeClient_ReleaseSegments_Call{Call: _e.mock.On("ReleaseSegments",
		append([]interface{}{ctx, in}, opts...)...)}
}

func (_c *MockQueryNodeClient_ReleaseSegments_Call) Run(run func(ctx context.Context, in *querypb.ReleaseSegmentsRequest, opts ...grpc.CallOption)) *MockQueryNodeClient_ReleaseSegments_Call {
	_c.Call.Run(func(args mock.Arguments) {
		variadicArgs := make([]grpc.CallOption, len(args)-2)
		for i, a := range args[2:] {
			if a != nil {
				variadicArgs[i] = a.(grpc.CallOption)
			}
		}
		run(args[0].(context.Context), args[1].(*querypb.ReleaseSegmentsRequest), variadicArgs...)
	})
	return _c
}

func (_c *MockQueryNodeClient_ReleaseSegments_Call) Return(_a0 *commonpb.Status, _a1 error) *MockQueryNodeClient_ReleaseSegments_Call {
	_c.Call.Return(_a0, _a1)
	return _c
}

func (_c *MockQueryNodeClient_ReleaseSegments_Call) RunAndReturn(run func(context.Context, *querypb.ReleaseSegmentsRequest, ...grpc.CallOption) (*commonpb.Status, error)) *MockQueryNodeClient_ReleaseSegments_Call {
	_c.Call.Return(run)
	return _c
}

// Search provides a mock function with given fields: ctx, in, opts
func (_m *MockQueryNodeClient) Search(ctx context.Context, in *querypb.SearchRequest, opts ...grpc.CallOption) (*internalpb.SearchResults, error) {
	_va := make([]interface{}, len(opts))
	for _i := range opts {
		_va[_i] = opts[_i]
	}
	var _ca []interface{}
	_ca = append(_ca, ctx, in)
	_ca = append(_ca, _va...)
	ret := _m.Called(_ca...)

	if len(ret) == 0 {
		panic("no return value specified for Search")
	}

	var r0 *internalpb.SearchResults
	var r1 error
	if rf, ok := ret.Get(0).(func(context.Context, *querypb.SearchRequest, ...grpc.CallOption) (*internalpb.SearchResults, error)); ok {
		return rf(ctx, in, opts...)
	}
	if rf, ok := ret.Get(0).(func(context.Context, *querypb.SearchRequest, ...grpc.CallOption) *internalpb.SearchResults); ok {
		r0 = rf(ctx, in, opts...)
	} else {
		if ret.Get(0) != nil {
			r0 = ret.Get(0).(*internalpb.SearchResults)
		}
	}

	if rf, ok := ret.Get(1).(func(context.Context, *querypb.SearchRequest, ...grpc.CallOption) error); ok {
		r1 = rf(ctx, in, opts...)
	} else {
		r1 = ret.Error(1)
	}

	return r0, r1
}

// MockQueryNodeClient_Search_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Search'
type MockQueryNodeClient_Search_Call struct {
	*mock.Call
}

// Search is a helper method to define mock.On call
//   - ctx context.Context
//   - in *querypb.SearchRequest
//   - opts ...grpc.CallOption
func (_e *MockQueryNodeClient_Expecter) Search(ctx interface{}, in interface{}, opts ...interface{}) *MockQueryNodeClient_Search_Call {
	return &MockQueryNodeClient_Search_Call{Call: _e.mock.On("Search",
		append([]interface{}{ctx, in}, opts...)...)}
}

func (_c *MockQueryNodeClient_Search_Call) Run(run func(ctx context.Context, in *querypb.SearchRequest, opts ...grpc.CallOption)) *MockQueryNodeClient_Search_Call {
	_c.Call.Run(func(args mock.Arguments) {
		variadicArgs := make([]grpc.CallOption, len(args)-2)
		for i, a := range args[2:] {
			if a != nil {
				variadicArgs[i] = a.(grpc.CallOption)
			}
		}
		run(args[0].(context.Context), args[1].(*querypb.SearchRequest), variadicArgs...)
	})
	return _c
}

func (_c *MockQueryNodeClient_Search_Call) Return(_a0 *internalpb.SearchResults, _a1 error) *MockQueryNodeClient_Search_Call {
	_c.Call.Return(_a0, _a1)
	return _c
}

func (_c *MockQueryNodeClient_Search_Call) RunAndReturn(run func(context.Context, *querypb.SearchRequest, ...grpc.CallOption) (*internalpb.SearchResults, error)) *MockQueryNodeClient_Search_Call {
	_c.Call.Return(run)
	return _c
}

// SearchSegments provides a mock function with given fields: ctx, in, opts
func (_m *MockQueryNodeClient) SearchSegments(ctx context.Context, in *querypb.SearchRequest, opts ...grpc.CallOption) (*internalpb.SearchResults, error) {
	_va := make([]interface{}, len(opts))
	for _i := range opts {
		_va[_i] = opts[_i]
	}
	var _ca []interface{}
	_ca = append(_ca, ctx, in)
	_ca = append(_ca, _va...)
	ret := _m.Called(_ca...)

	if len(ret) == 0 {
		panic("no return value specified for SearchSegments")
	}

	var r0 *internalpb.SearchResults
	var r1 error
	if rf, ok := ret.Get(0).(func(context.Context, *querypb.SearchRequest, ...grpc.CallOption) (*internalpb.SearchResults, error)); ok {
		return rf(ctx, in, opts...)
	}
	if rf, ok := ret.Get(0).(func(context.Context, *querypb.SearchRequest, ...grpc.CallOption) *internalpb.SearchResults); ok {
		r0 = rf(ctx, in, opts...)
	} else {
		if ret.Get(0) != nil {
			r0 = ret.Get(0).(*internalpb.SearchResults)
		}
	}

	if rf, ok := ret.Get(1).(func(context.Context, *querypb.SearchRequest, ...grpc.CallOption) error); ok {
		r1 = rf(ctx, in, opts...)
	} else {
		r1 = ret.Error(1)
	}

	return r0, r1
}

// MockQueryNodeClient_SearchSegments_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SearchSegments'
type MockQueryNodeClient_SearchSegments_Call struct {
	*mock.Call
}

// SearchSegments is a helper method to define mock.On call
//   - ctx context.Context
//   - in *querypb.SearchRequest
//   - opts ...grpc.CallOption
func (_e *MockQueryNodeClient_Expecter) SearchSegments(ctx interface{}, in interface{}, opts ...interface{}) *MockQueryNodeClient_SearchSegments_Call {
	return &MockQueryNodeClient_SearchSegments_Call{Call: _e.mock.On("SearchSegments",
		append([]interface{}{ctx, in}, opts...)...)}
}

func (_c *MockQueryNodeClient_SearchSegments_Call) Run(run func(ctx context.Context, in *querypb.SearchRequest, opts ...grpc.CallOption)) *MockQueryNodeClient_SearchSegments_Call {
	_c.Call.Run(func(args mock.Arguments) {
		variadicArgs := make([]grpc.CallOption, len(args)-2)
		for i, a := range args[2:] {
			if a != nil {
				variadicArgs[i] = a.(grpc.CallOption)
			}
		}
		run(args[0].(context.Context), args[1].(*querypb.SearchRequest), variadicArgs...)
	})
	return _c
}

func (_c *MockQueryNodeClient_SearchSegments_Call) Return(_a0 *internalpb.SearchResults, _a1 error) *MockQueryNodeClient_SearchSegments_Call {
	_c.Call.Return(_a0, _a1)
	return _c
}

func (_c *MockQueryNodeClient_SearchSegments_Call) RunAndReturn(run func(context.Context, *querypb.SearchRequest, ...grpc.CallOption) (*internalpb.SearchResults, error)) *MockQueryNodeClient_SearchSegments_Call {
	_c.Call.Return(run)
	return _c
}

// ShowConfigurations provides a mock function with given fields: ctx, in, opts
func (_m *MockQueryNodeClient) ShowConfigurations(ctx context.Context, in *internalpb.ShowConfigurationsRequest, opts ...grpc.CallOption) (*internalpb.ShowConfigurationsResponse, error) {
	_va := make([]interface{}, len(opts))
	for _i := range opts {
		_va[_i] = opts[_i]
	}
	var _ca []interface{}
	_ca = append(_ca, ctx, in)
	_ca = append(_ca, _va...)
	ret := _m.Called(_ca...)

	if len(ret) == 0 {
		panic("no return value specified for ShowConfigurations")
	}

	var r0 *internalpb.ShowConfigurationsResponse
	var r1 error
	if rf, ok := ret.Get(0).(func(context.Context, *internalpb.ShowConfigurationsRequest, ...grpc.CallOption) (*internalpb.ShowConfigurationsResponse, error)); ok {
		return rf(ctx, in, opts...)
	}
	if rf, ok := ret.Get(0).(func(context.Context, *internalpb.ShowConfigurationsRequest, ...grpc.CallOption) *internalpb.ShowConfigurationsResponse); ok {
		r0 = rf(ctx, in, opts...)
	} else {
		if ret.Get(0) != nil {
			r0 = ret.Get(0).(*internalpb.ShowConfigurationsResponse)
		}
	}

	if rf, ok := ret.Get(1).(func(context.Context, *internalpb.ShowConfigurationsRequest, ...grpc.CallOption) error); ok {
		r1 = rf(ctx, in, opts...)
	} else {
		r1 = ret.Error(1)
	}

	return r0, r1
}

// MockQueryNodeClient_ShowConfigurations_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'ShowConfigurations'
type MockQueryNodeClient_ShowConfigurations_Call struct {
	*mock.Call
}

// ShowConfigurations is a helper method to define mock.On call
//   - ctx context.Context
//   - in *internalpb.ShowConfigurationsRequest
//   - opts ...grpc.CallOption
func (_e *MockQueryNodeClient_Expecter) ShowConfigurations(ctx interface{}, in interface{}, opts ...interface{}) *MockQueryNodeClient_ShowConfigurations_Call {
	return &MockQueryNodeClient_ShowConfigurations_Call{Call: _e.mock.On("ShowConfigurations",
		append([]interface{}{ctx, in}, opts...)...)}
}

func (_c *MockQueryNodeClient_ShowConfigurations_Call) Run(run func(ctx context.Context, in *internalpb.ShowConfigurationsRequest, opts ...grpc.CallOption)) *MockQueryNodeClient_ShowConfigurations_Call {
	_c.Call.Run(func(args mock.Arguments) {
		variadicArgs := make([]grpc.CallOption, len(args)-2)
		for i, a := range args[2:] {
			if a != nil {
				variadicArgs[i] = a.(grpc.CallOption)
			}
		}
		run(args[0].(context.Context), args[1].(*internalpb.ShowConfigurationsRequest), variadicArgs...)
	})
	return _c
}

func (_c *MockQueryNodeClient_ShowConfigurations_Call) Return(_a0 *internalpb.ShowConfigurationsResponse, _a1 error) *MockQueryNodeClient_ShowConfigurations_Call {
	_c.Call.Return(_a0, _a1)
	return _c
}

func (_c *MockQueryNodeClient_ShowConfigurations_Call) RunAndReturn(run func(context.Context, *internalpb.ShowConfigurationsRequest, ...grpc.CallOption) (*internalpb.ShowConfigurationsResponse, error)) *MockQueryNodeClient_ShowConfigurations_Call {
	_c.Call.Return(run)
	return _c
}

// SyncDistribution provides a mock function with given fields: ctx, in, opts
func (_m *MockQueryNodeClient) SyncDistribution(ctx context.Context, in *querypb.SyncDistributionRequest, opts ...grpc.CallOption) (*commonpb.Status, error) {
	_va := make([]interface{}, len(opts))
	for _i := range opts {
		_va[_i] = opts[_i]
	}
	var _ca []interface{}
	_ca = append(_ca, ctx, in)
	_ca = append(_ca, _va...)
	ret := _m.Called(_ca...)

	if len(ret) == 0 {
		panic("no return value specified for SyncDistribution")
	}

	var r0 *commonpb.Status
	var r1 error
	if rf, ok := ret.Get(0).(func(context.Context, *querypb.SyncDistributionRequest, ...grpc.CallOption) (*commonpb.Status, error)); ok {
		return rf(ctx, in, opts...)
	}
	if rf, ok := ret.Get(0).(func(context.Context, *querypb.SyncDistributionRequest, ...grpc.CallOption) *commonpb.Status); ok {
		r0 = rf(ctx, in, opts...)
	} else {
		if ret.Get(0) != nil {
			r0 = ret.Get(0).(*commonpb.Status)
		}
	}

	if rf, ok := ret.Get(1).(func(context.Context, *querypb.SyncDistributionRequest, ...grpc.CallOption) error); ok {
		r1 = rf(ctx, in, opts...)
	} else {
		r1 = ret.Error(1)
	}

	return r0, r1
}

// MockQueryNodeClient_SyncDistribution_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SyncDistribution'
type MockQueryNodeClient_SyncDistribution_Call struct {
	*mock.Call
}

// SyncDistribution is a helper method to define mock.On call
//   - ctx context.Context
//   - in *querypb.SyncDistributionRequest
//   - opts ...grpc.CallOption
func (_e *MockQueryNodeClient_Expecter) SyncDistribution(ctx interface{}, in interface{}, opts ...interface{}) *MockQueryNodeClient_SyncDistribution_Call {
	return &MockQueryNodeClient_SyncDistribution_Call{Call: _e.mock.On("SyncDistribution",
		append([]interface{}{ctx, in}, opts...)...)}
}

func (_c *MockQueryNodeClient_SyncDistribution_Call) Run(run func(ctx context.Context, in *querypb.SyncDistributionRequest, opts ...grpc.CallOption)) *MockQueryNodeClient_SyncDistribution_Call {
	_c.Call.Run(func(args mock.Arguments) {
		variadicArgs := make([]grpc.CallOption, len(args)-2)
		for i, a := range args[2:] {
			if a != nil {
				variadicArgs[i] = a.(grpc.CallOption)
			}
		}
		run(args[0].(context.Context), args[1].(*querypb.SyncDistributionRequest), variadicArgs...)
	})
	return _c
}

func (_c *MockQueryNodeClient_SyncDistribution_Call) Return(_a0 *commonpb.Status, _a1 error) *MockQueryNodeClient_SyncDistribution_Call {
	_c.Call.Return(_a0, _a1)
	return _c
}

func (_c *MockQueryNodeClient_SyncDistribution_Call) RunAndReturn(run func(context.Context, *querypb.SyncDistributionRequest, ...grpc.CallOption) (*commonpb.Status, error)) *MockQueryNodeClient_SyncDistribution_Call {
	_c.Call.Return(run)
	return _c
}

// SyncReplicaSegments provides a mock function with given fields: ctx, in, opts
func (_m *MockQueryNodeClient) SyncReplicaSegments(ctx context.Context, in *querypb.SyncReplicaSegmentsRequest, opts ...grpc.CallOption) (*commonpb.Status, error) {
	_va := make([]interface{}, len(opts))
	for _i := range opts {
		_va[_i] = opts[_i]
	}
	var _ca []interface{}
	_ca = append(_ca, ctx, in)
	_ca = append(_ca, _va...)
	ret := _m.Called(_ca...)

	if len(ret) == 0 {
		panic("no return value specified for SyncReplicaSegments")
	}

	var r0 *commonpb.Status
	var r1 error
	if rf, ok := ret.Get(0).(func(context.Context, *querypb.SyncReplicaSegmentsRequest, ...grpc.CallOption) (*commonpb.Status, error)); ok {
		return rf(ctx, in, opts...)
	}
	if rf, ok := ret.Get(0).(func(context.Context, *querypb.SyncReplicaSegmentsRequest, ...grpc.CallOption) *commonpb.Status); ok {
		r0 = rf(ctx, in, opts...)
	} else {
		if ret.Get(0) != nil {
			r0 = ret.Get(0).(*commonpb.Status)
		}
	}

	if rf, ok := ret.Get(1).(func(context.Context, *querypb.SyncReplicaSegmentsRequest, ...grpc.CallOption) error); ok {
		r1 = rf(ctx, in, opts...)
	} else {
		r1 = ret.Error(1)
	}

	return r0, r1
}

// MockQueryNodeClient_SyncReplicaSegments_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SyncReplicaSegments'
type MockQueryNodeClient_SyncReplicaSegments_Call struct {
	*mock.Call
}

// SyncReplicaSegments is a helper method to define mock.On call
//   - ctx context.Context
//   - in *querypb.SyncReplicaSegmentsRequest
//   - opts ...grpc.CallOption
func (_e *MockQueryNodeClient_Expecter) SyncReplicaSegments(ctx interface{}, in interface{}, opts ...interface{}) *MockQueryNodeClient_SyncReplicaSegments_Call {
	return &MockQueryNodeClient_SyncReplicaSegments_Call{Call: _e.mock.On("SyncReplicaSegments",
		append([]interface{}{ctx, in}, opts...)...)}
}

func (_c *MockQueryNodeClient_SyncReplicaSegments_Call) Run(run func(ctx context.Context, in *querypb.SyncReplicaSegmentsRequest, opts ...grpc.CallOption)) *MockQueryNodeClient_SyncReplicaSegments_Call {
	_c.Call.Run(func(args mock.Arguments) {
		variadicArgs := make([]grpc.CallOption, len(args)-2)
		for i, a := range args[2:] {
			if a != nil {
				variadicArgs[i] = a.(grpc.CallOption)
			}
		}
		run(args[0].(context.Context), args[1].(*querypb.SyncReplicaSegmentsRequest), variadicArgs...)
	})
	return _c
}

func (_c *MockQueryNodeClient_SyncReplicaSegments_Call) Return(_a0 *commonpb.Status, _a1 error) *MockQueryNodeClient_SyncReplicaSegments_Call {
	_c.Call.Return(_a0, _a1)
	return _c
}

func (_c *MockQueryNodeClient_SyncReplicaSegments_Call) RunAndReturn(run func(context.Context, *querypb.SyncReplicaSegmentsRequest, ...grpc.CallOption) (*commonpb.Status, error)) *MockQueryNodeClient_SyncReplicaSegments_Call {
	_c.Call.Return(run)
	return _c
}

// UnsubDmChannel provides a mock function with given fields: ctx, in, opts
func (_m *MockQueryNodeClient) UnsubDmChannel(ctx context.Context, in *querypb.UnsubDmChannelRequest, opts ...grpc.CallOption) (*commonpb.Status, error) {
	_va := make([]interface{}, len(opts))
	for _i := range opts {
		_va[_i] = opts[_i]
	}
	var _ca []interface{}
	_ca = append(_ca, ctx, in)
	_ca = append(_ca, _va...)
	ret := _m.Called(_ca...)

	if len(ret) == 0 {
		panic("no return value specified for UnsubDmChannel")
	}

	var r0 *commonpb.Status
	var r1 error
	if rf, ok := ret.Get(0).(func(context.Context, *querypb.UnsubDmChannelRequest, ...grpc.CallOption) (*commonpb.Status, error)); ok {
		return rf(ctx, in, opts...)
	}
	if rf, ok := ret.Get(0).(func(context.Context, *querypb.UnsubDmChannelRequest, ...grpc.CallOption) *commonpb.Status); ok {
		r0 = rf(ctx, in, opts...)
	} else {
		if ret.Get(0) != nil {
			r0 = ret.Get(0).(*commonpb.Status)
		}
	}

	if rf, ok := ret.Get(1).(func(context.Context, *querypb.UnsubDmChannelRequest, ...grpc.CallOption) error); ok {
		r1 = rf(ctx, in, opts...)
	} else {
		r1 = ret.Error(1)
	}

	return r0, r1
}

// MockQueryNodeClient_UnsubDmChannel_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'UnsubDmChannel'
type MockQueryNodeClient_UnsubDmChannel_Call struct {
	*mock.Call
}

// UnsubDmChannel is a helper method to define mock.On call
//   - ctx context.Context
//   - in *querypb.UnsubDmChannelRequest
//   - opts ...grpc.CallOption
func (_e *MockQueryNodeClient_Expecter) UnsubDmChannel(ctx interface{}, in interface{}, opts ...interface{}) *MockQueryNodeClient_UnsubDmChannel_Call {
	return &MockQueryNodeClient_UnsubDmChannel_Call{Call: _e.mock.On("UnsubDmChannel",
		append([]interface{}{ctx, in}, opts...)...)}
}

func (_c *MockQueryNodeClient_UnsubDmChannel_Call) Run(run func(ctx context.Context, in *querypb.UnsubDmChannelRequest, opts ...grpc.CallOption)) *MockQueryNodeClient_UnsubDmChannel_Call {
	_c.Call.Run(func(args mock.Arguments) {
		variadicArgs := make([]grpc.CallOption, len(args)-2)
		for i, a := range args[2:] {
			if a != nil {
				variadicArgs[i] = a.(grpc.CallOption)
			}
		}
		run(args[0].(context.Context), args[1].(*querypb.UnsubDmChannelRequest), variadicArgs...)
	})
	return _c
}

func (_c *MockQueryNodeClient_UnsubDmChannel_Call) Return(_a0 *commonpb.Status, _a1 error) *MockQueryNodeClient_UnsubDmChannel_Call {
	_c.Call.Return(_a0, _a1)
	return _c
}

func (_c *MockQueryNodeClient_UnsubDmChannel_Call) RunAndReturn(run func(context.Context, *querypb.UnsubDmChannelRequest, ...grpc.CallOption) (*commonpb.Status, error)) *MockQueryNodeClient_UnsubDmChannel_Call {
	_c.Call.Return(run)
	return _c
}

// WatchDmChannels provides a mock function with given fields: ctx, in, opts
func (_m *MockQueryNodeClient) WatchDmChannels(ctx context.Context, in *querypb.WatchDmChannelsRequest, opts ...grpc.CallOption) (*commonpb.Status, error) {
	_va := make([]interface{}, len(opts))
	for _i := range opts {
		_va[_i] = opts[_i]
	}
	var _ca []interface{}
	_ca = append(_ca, ctx, in)
	_ca = append(_ca, _va...)
	ret := _m.Called(_ca...)

	if len(ret) == 0 {
		panic("no return value specified for WatchDmChannels")
	}

	var r0 *commonpb.Status
	var r1 error
	if rf, ok := ret.Get(0).(func(context.Context, *querypb.WatchDmChannelsRequest, ...grpc.CallOption) (*commonpb.Status, error)); ok {
		return rf(ctx, in, opts...)
	}
	if rf, ok := ret.Get(0).(func(context.Context, *querypb.WatchDmChannelsRequest, ...grpc.CallOption) *commonpb.Status); ok {
		r0 = rf(ctx, in, opts...)
	} else {
		if ret.Get(0) != nil {
			r0 = ret.Get(0).(*commonpb.Status)
		}
	}

	if rf, ok := ret.Get(1).(func(context.Context, *querypb.WatchDmChannelsRequest, ...grpc.CallOption) error); ok {
		r1 = rf(ctx, in, opts...)
	} else {
		r1 = ret.Error(1)
	}

	return r0, r1
}

// MockQueryNodeClient_WatchDmChannels_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'WatchDmChannels'
type MockQueryNodeClient_WatchDmChannels_Call struct {
	*mock.Call
}

// WatchDmChannels is a helper method to define mock.On call
//   - ctx context.Context
//   - in *querypb.WatchDmChannelsRequest
//   - opts ...grpc.CallOption
func (_e *MockQueryNodeClient_Expecter) WatchDmChannels(ctx interface{}, in interface{}, opts ...interface{}) *MockQueryNodeClient_WatchDmChannels_Call {
	return &MockQueryNodeClient_WatchDmChannels_Call{Call: _e.mock.On("WatchDmChannels",
		append([]interface{}{ctx, in}, opts...)...)}
}

func (_c *MockQueryNodeClient_WatchDmChannels_Call) Run(run func(ctx context.Context, in *querypb.WatchDmChannelsRequest, opts ...grpc.CallOption)) *MockQueryNodeClient_WatchDmChannels_Call {
	_c.Call.Run(func(args mock.Arguments) {
		variadicArgs := make([]grpc.CallOption, len(args)-2)
		for i, a := range args[2:] {
			if a != nil {
				variadicArgs[i] = a.(grpc.CallOption)
			}
		}
		run(args[0].(context.Context), args[1].(*querypb.WatchDmChannelsRequest), variadicArgs...)
	})
	return _c
}

func (_c *MockQueryNodeClient_WatchDmChannels_Call) Return(_a0 *commonpb.Status, _a1 error) *MockQueryNodeClient_WatchDmChannels_Call {
	_c.Call.Return(_a0, _a1)
	return _c
}

func (_c *MockQueryNodeClient_WatchDmChannels_Call) RunAndReturn(run func(context.Context, *querypb.WatchDmChannelsRequest, ...grpc.CallOption) (*commonpb.Status, error)) *MockQueryNodeClient_WatchDmChannels_Call {
	_c.Call.Return(run)
	return _c
}

// NewMockQueryNodeClient creates a new instance of MockQueryNodeClient. 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 NewMockQueryNodeClient(t interface {
	mock.TestingT
	Cleanup(func())
}) *MockQueryNodeClient {
	mock := &MockQueryNodeClient{}
	mock.Mock.Test(t)

	t.Cleanup(func() { mock.AssertExpectations(t) })

	return mock
}