Add unittest for distributed/querynode (#9608)

Signed-off-by: yhmo <yihua.mo@zilliz.com>

internal/distributed/querynode/service.go

@@ -45,7 +45,7 @@ import (
 type UniqueID = typeutil.UniqueID
 
 type Server struct {
-	querynode   qn.Base
+	querynode   types.QueryNodeComponent
 	wg          sync.WaitGroup
 	ctx         context.Context
 	cancel      context.CancelFunc
@@ -259,12 +259,6 @@ func (s *Server) SetIndexCoord(indexCoord types.IndexCoord) error {
 	return s.querynode.SetIndexCoord(indexCoord)
 }
 
-// SetClient sets the IndexNode's instance.
-func (s *Server) SetClient(queryNodeClient qn.Base) error {
-	s.querynode = queryNodeClient
-	return nil
-}
-
 func (s *Server) GetTimeTickChannel(ctx context.Context, req *internalpb.GetTimeTickChannelRequest) (*milvuspb.StringResponse, error) {
 	return s.querynode.GetTimeTickChannel(ctx)
 }

internal/distributed/querynode/service_test.go

@@ -33,25 +33,29 @@ type MockQueryNode struct {
 	states     *internalpb.ComponentStates
 	status     *commonpb.Status
 	err        error
+	initErr    error
+	startErr   error
+	regErr     error
+	stopErr    error
 	strResp    *milvuspb.StringResponse
 	infoResp   *querypb.GetSegmentInfoResponse
 	metricResp *milvuspb.GetMetricsResponse
 }
 
 func (m *MockQueryNode) Init() error {
-	return m.err
+	return m.initErr
 }
 
 func (m *MockQueryNode) Start() error {
-	return m.err
+	return m.startErr
 }
 
 func (m *MockQueryNode) Stop() error {
-	return m.err
+	return m.stopErr
 }
 
 func (m *MockQueryNode) Register() error {
-	return m.err
+	return m.regErr
 }
 
 func (m *MockQueryNode) GetComponentStates(ctx context.Context) (*internalpb.ComponentStates, error) {
@@ -194,8 +198,7 @@ func Test_NewServer(t *testing.T) {
 		infoResp:   &querypb.GetSegmentInfoResponse{Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}},
 		metricResp: &milvuspb.GetMetricsResponse{Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}},
 	}
-	err = qns.SetClient(mqn)
-	assert.Nil(t, err)
+	qns.querynode = mqn
 
 	t.Run("Run", func(t *testing.T) {
 		qns.rootCoord = &MockRootCoord{}
@@ -301,19 +304,37 @@ func Test_Run(t *testing.T) {
 	assert.Nil(t, err)
 	assert.NotNil(t, qns)
 
-	qns.rootCoord = &MockRootCoord{initErr: errors.New("Failed")}
+	qns.querynode = &MockQueryNode{}
 	qns.indexCoord = &MockIndexCoord{}
+	qns.rootCoord = &MockRootCoord{initErr: errors.New("Failed")}
 	assert.Panics(t, func() { err = qns.Run() })
 
 	qns.rootCoord = &MockRootCoord{startErr: errors.New("Failed")}
-	qns.indexCoord = &MockIndexCoord{}
 	assert.Panics(t, func() { err = qns.Run() })
 
+	qns.querynode = &MockQueryNode{}
 	qns.rootCoord = &MockRootCoord{}
 	qns.indexCoord = &MockIndexCoord{initErr: errors.New("Failed")}
 	assert.Panics(t, func() { err = qns.Run() })
 
-	qns.rootCoord = &MockRootCoord{}
 	qns.indexCoord = &MockIndexCoord{startErr: errors.New("Failed")}
 	assert.Panics(t, func() { err = qns.Run() })
+
+	qns.indexCoord = &MockIndexCoord{}
+	qns.rootCoord = &MockRootCoord{}
+	qns.querynode = &MockQueryNode{initErr: errors.New("Failed")}
+	err = qns.Run()
+	assert.Error(t, err)
+
+	qns.querynode = &MockQueryNode{startErr: errors.New("Failed")}
+	err = qns.Run()
+	assert.Error(t, err)
+
+	qns.querynode = &MockQueryNode{regErr: errors.New("Failed")}
+	err = qns.Run()
+	assert.Error(t, err)
+
+	qns.querynode = &MockQueryNode{stopErr: errors.New("Failed")}
+	err = qns.Stop()
+	assert.Error(t, err)
 }

internal/querynode/query_node.go

@@ -45,19 +45,11 @@ import (
 	"github.com/milvus-io/milvus/internal/util/typeutil"
 )
 
-type Base interface {
-	types.QueryNode
-
-	UpdateStateCode(code internalpb.StateCode)
-	SetRootCoord(rc types.RootCoord) error
-	SetIndexCoord(index types.IndexCoord) error
-}
-
 // make sure QueryNode implements types.QueryNode
 var _ types.QueryNode = (*QueryNode)(nil)
 
-// make sure QueryNode implements Base
-var _ Base = (*QueryNode)(nil)
+// make sure QueryNode implements types.QueryNodeComponent
+var _ types.QueryNodeComponent = (*QueryNode)(nil)
 
 // QueryNode communicates with outside services and union all
 // services in querynode package.

internal/types/types.go

@@ -533,6 +533,33 @@ type QueryNode interface {
 	GetMetrics(ctx context.Context, req *milvuspb.GetMetricsRequest) (*milvuspb.GetMetricsResponse, error)
 }
 
+// QueryNodeComponent is used by grpc server of QueryNode
+type QueryNodeComponent interface {
+	QueryNode
+
+	// UpdateStateCode updates state code for QueryNode
+	//  `stateCode` is current statement of this query node, indicating whether it's healthy.
+	UpdateStateCode(code internalpb.StateCode)
+
+	// SetRootCoord set RootCoord for QueryNode
+	// `rootCoord` is a client of root coordinator. Pass to segmentLoader.
+	//
+	// Return a generic error in status:
+	//     If the rootCoord is nil.
+	// Return nil in status:
+	//     The rootCoord is not nil.
+	SetRootCoord(rootCoord RootCoord) error
+
+	// SetIndexCoord set IndexCoord for QueryNode
+	//  `indexCoord` is a client of index coordinator. Pass to segmentLoader.
+	//
+	// Return a generic error in status:
+	//     If the indexCoord is nil.
+	// Return nil in status:
+	//     The indexCoord is not nil.
+	SetIndexCoord(indexCoord IndexCoord) error
+}
+
 // QueryCoord is the interface `querycoord` package implements
 type QueryCoord interface {
 	Component