Add unittest for distributed/querycoord (#9567)

Signed-off-by: yhmo <yihua.mo@zilliz.com>
2026-02-02 01:06:41 +08:00 · 2021-10-11 10:02:25 +08:00 · 2021-10-11 10:02:25 +08:00 · 1277bb659a
commit 1277bb659a
parent 235d736a49
4 changed files with 458 additions and 20 deletions
--- a/internal/distributed/datacoord/client/client.go
+++ b/internal/distributed/datacoord/client/client.go
@ -22,6 +22,7 @@ import (
 	grpc_opentracing "github.com/grpc-ecosystem/go-grpc-middleware/tracing/opentracing"
 	"github.com/milvus-io/milvus/internal/log"
 	"github.com/milvus-io/milvus/internal/proto/milvuspb"
+	"github.com/milvus-io/milvus/internal/types"
 	"github.com/milvus-io/milvus/internal/util/retry"
 	"github.com/milvus-io/milvus/internal/util/sessionutil"
 	"github.com/milvus-io/milvus/internal/util/trace"
@ -35,6 +36,16 @@ import (
 	"github.com/milvus-io/milvus/internal/proto/internalpb"
 )

+// Base is a base class abstracted from components.
+type Base interface {
+	types.DataCoord
+
+	Init() error
+	Start() error
+	Stop() error
+	Register() error
+}
+
 type Client struct {
 	ctx    context.Context
 	cancel context.CancelFunc
--- a/internal/distributed/querycoord/service.go
+++ b/internal/distributed/querycoord/service.go
@ -45,12 +45,12 @@ type Server struct {

 	grpcErrChan chan error

-	queryCoord *qc.QueryCoord
+	queryCoord types.QueryCoordComponent

 	msFactory msgstream.Factory

-	dataCoord *dsc.Client
-	rootCoord *rcc.GrpcClient
+	dataCoord dsc.Base
+	rootCoord rcc.Base

 	closer io.Closer
 }
@ -101,36 +101,40 @@ func (s *Server) init() error {
 	s.wg.Add(1)
 	go s.startGrpcLoop(Params.Port)
 	// wait for grpc server loop start
-	if err := <-s.grpcErrChan; err != nil {
+	err := <-s.grpcErrChan
+	if err != nil {
 		return err
 	}

 	// --- Master Server Client ---
 	log.Debug("QueryCoord try to new RootCoord client", zap.Any("RootCoordAddress", Params.RootCoordAddress))
-	rootCoord, err := rcc.NewClient(s.loopCtx, qc.Params.MetaRootPath, qc.Params.EtcdEndpoints)
-	if err != nil {
-		log.Debug("QueryCoord try to new RootCoord client failed", zap.Error(err))
-		panic(err)
+
+	if s.rootCoord == nil {
+		s.rootCoord, err = rcc.NewClient(s.loopCtx, qc.Params.MetaRootPath, qc.Params.EtcdEndpoints)
+		if err != nil {
+			log.Debug("QueryCoord try to new RootCoord client failed", zap.Error(err))
+			panic(err)
+		}
 	}

-	if err = rootCoord.Init(); err != nil {
+	if err = s.rootCoord.Init(); err != nil {
 		log.Debug("QueryCoord RootCoordClient Init failed", zap.Error(err))
 		panic(err)
 	}

-	if err = rootCoord.Start(); err != nil {
+	if err = s.rootCoord.Start(); err != nil {
 		log.Debug("QueryCoord RootCoordClient Start failed", zap.Error(err))
 		panic(err)
 	}
 	// wait for master init or healthy
 	log.Debug("QueryCoord try to wait for RootCoord ready")
-	err = funcutil.WaitForComponentInitOrHealthy(s.loopCtx, rootCoord, "RootCoord", 1000000, time.Millisecond*200)
+	err = funcutil.WaitForComponentInitOrHealthy(s.loopCtx, s.rootCoord, "RootCoord", 1000000, time.Millisecond*200)
 	if err != nil {
 		log.Debug("QueryCoord wait for RootCoord ready failed", zap.Error(err))
 		panic(err)
 	}

-	if err := s.SetRootCoord(rootCoord); err != nil {
+	if err := s.SetRootCoord(s.rootCoord); err != nil {
 		panic(err)
 	}
 	log.Debug("QueryCoord report RootCoord ready")
@ -138,26 +142,29 @@ func (s *Server) init() error {
 	// --- Data service client ---
 	log.Debug("QueryCoord try to new DataCoord client", zap.Any("DataCoordAddress", Params.DataCoordAddress))

-	dataCoord, err := dsc.NewClient(s.loopCtx, qc.Params.MetaRootPath, qc.Params.EtcdEndpoints)
-	if err != nil {
-		log.Debug("QueryCoord try to new DataCoord client failed", zap.Error(err))
-		panic(err)
+	if s.dataCoord == nil {
+		s.dataCoord, err = dsc.NewClient(s.loopCtx, qc.Params.MetaRootPath, qc.Params.EtcdEndpoints)
+		if err != nil {
+			log.Debug("QueryCoord try to new DataCoord client failed", zap.Error(err))
+			panic(err)
+		}
 	}
-	if err = dataCoord.Init(); err != nil {
+
+	if err = s.dataCoord.Init(); err != nil {
 		log.Debug("QueryCoord DataCoordClient Init failed", zap.Error(err))
 		panic(err)
 	}
-	if err = dataCoord.Start(); err != nil {
+	if err = s.dataCoord.Start(); err != nil {
 		log.Debug("QueryCoord DataCoordClient Start failed", zap.Error(err))
 		panic(err)
 	}
 	log.Debug("QueryCoord try to wait for DataCoord ready")
-	err = funcutil.WaitForComponentInitOrHealthy(s.loopCtx, dataCoord, "DataCoord", 1000000, time.Millisecond*200)
+	err = funcutil.WaitForComponentInitOrHealthy(s.loopCtx, s.dataCoord, "DataCoord", 1000000, time.Millisecond*200)
 	if err != nil {
 		log.Debug("QueryCoord wait for DataCoord ready failed", zap.Error(err))
 		panic(err)
 	}
-	if err := s.SetDataCoord(dataCoord); err != nil {
+	if err := s.SetDataCoord(s.dataCoord); err != nil {
 		panic(err)
 	}
 	log.Debug("QueryCoord report DataCoord ready")
--- a/internal/distributed/querycoord/service_test.go
+++ b/internal/distributed/querycoord/service_test.go
@ -0,0 +1,405 @@
+// Copyright (C) 2019-2020 Zilliz. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software distributed under the License
+// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+// or implied. See the License for the specific language governing permissions and limitations under the License.
+
+package grpcquerycoord
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	dcc "github.com/milvus-io/milvus/internal/distributed/datacoord/client"
+	rcc "github.com/milvus-io/milvus/internal/distributed/rootcoord/client"
+	"github.com/milvus-io/milvus/internal/log"
+	"github.com/milvus-io/milvus/internal/types"
+
+	"github.com/milvus-io/milvus/internal/proto/commonpb"
+	"github.com/milvus-io/milvus/internal/proto/internalpb"
+	"github.com/milvus-io/milvus/internal/proto/milvuspb"
+	"github.com/milvus-io/milvus/internal/proto/querypb"
+	"github.com/stretchr/testify/assert"
+)
+
+///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
+type MockQueryCoord struct {
+	states       *internalpb.ComponentStates
+	status       *commonpb.Status
+	err          error
+	initErr      error
+	startErr     error
+	stopErr      error
+	regErr       error
+	strResp      *milvuspb.StringResponse
+	showcolResp  *querypb.ShowCollectionsResponse
+	showpartResp *querypb.ShowPartitionsResponse
+	partResp     *querypb.GetPartitionStatesResponse
+	channelResp  *querypb.CreateQueryChannelResponse
+	infoResp     *querypb.GetSegmentInfoResponse
+	metricResp   *milvuspb.GetMetricsResponse
+}
+
+func (m *MockQueryCoord) Init() error {
+	return m.initErr
+}
+
+func (m *MockQueryCoord) Start() error {
+	return m.startErr
+}
+
+func (m *MockQueryCoord) Stop() error {
+	return m.stopErr
+}
+
+func (m *MockQueryCoord) Register() error {
+	log.Debug("MockQueryCoord::Register")
+	return m.regErr
+}
+
+func (m *MockQueryCoord) UpdateStateCode(code internalpb.StateCode) {
+}
+
+func (m *MockQueryCoord) SetRootCoord(types.RootCoord) {
+}
+
+func (m *MockQueryCoord) SetDataCoord(types.DataCoord) {
+}
+
+func (m *MockQueryCoord) GetComponentStates(ctx context.Context) (*internalpb.ComponentStates, error) {
+	log.Debug("MockQueryCoord::WaitForComponentStates")
+	return m.states, m.err
+}
+
+func (m *MockQueryCoord) GetTimeTickChannel(ctx context.Context) (*milvuspb.StringResponse, error) {
+	return m.strResp, m.err
+}
+
+func (m *MockQueryCoord) GetStatisticsChannel(ctx context.Context) (*milvuspb.StringResponse, error) {
+	return m.strResp, m.err
+}
+
+func (m *MockQueryCoord) ShowCollections(ctx context.Context, req *querypb.ShowCollectionsRequest) (*querypb.ShowCollectionsResponse, error) {
+	return m.showcolResp, m.err
+}
+
+func (m *MockQueryCoord) LoadCollection(ctx context.Context, req *querypb.LoadCollectionRequest) (*commonpb.Status, error) {
+	return m.status, m.err
+}
+
+func (m *MockQueryCoord) ReleaseCollection(ctx context.Context, req *querypb.ReleaseCollectionRequest) (*commonpb.Status, error) {
+	return m.status, m.err
+}
+
+func (m *MockQueryCoord) ShowPartitions(ctx context.Context, req *querypb.ShowPartitionsRequest) (*querypb.ShowPartitionsResponse, error) {
+	return m.showpartResp, m.err
+}
+
+func (m *MockQueryCoord) GetPartitionStates(ctx context.Context, req *querypb.GetPartitionStatesRequest) (*querypb.GetPartitionStatesResponse, error) {
+	return m.partResp, m.err
+}
+
+func (m *MockQueryCoord) LoadPartitions(ctx context.Context, req *querypb.LoadPartitionsRequest) (*commonpb.Status, error) {
+	return m.status, m.err
+}
+
+func (m *MockQueryCoord) ReleasePartitions(ctx context.Context, req *querypb.ReleasePartitionsRequest) (*commonpb.Status, error) {
+	return m.status, m.err
+}
+
+func (m *MockQueryCoord) CreateQueryChannel(ctx context.Context, req *querypb.CreateQueryChannelRequest) (*querypb.CreateQueryChannelResponse, error) {
+	return m.channelResp, m.err
+}
+
+func (m *MockQueryCoord) GetSegmentInfo(ctx context.Context, req *querypb.GetSegmentInfoRequest) (*querypb.GetSegmentInfoResponse, error) {
+	return m.infoResp, m.err
+}
+
+func (m *MockQueryCoord) GetMetrics(ctx context.Context, req *milvuspb.GetMetricsRequest) (*milvuspb.GetMetricsResponse, error) {
+	return m.metricResp, m.err
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
+type MockRootCoord struct {
+	rcc.Base
+	initErr  error
+	startErr error
+	regErr   error
+	stopErr  error
+	stateErr commonpb.ErrorCode
+}
+
+func (m *MockRootCoord) Init() error {
+	return m.initErr
+}
+
+func (m *MockRootCoord) Start() error {
+	return m.startErr
+}
+
+func (m *MockRootCoord) Stop() error {
+	return m.stopErr
+}
+
+func (m *MockRootCoord) Register() error {
+	return m.regErr
+}
+
+func (m *MockRootCoord) GetComponentStates(ctx context.Context) (*internalpb.ComponentStates, error) {
+	return &internalpb.ComponentStates{
+		State:  &internalpb.ComponentInfo{StateCode: internalpb.StateCode_Healthy},
+		Status: &commonpb.Status{ErrorCode: m.stateErr},
+	}, nil
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
+type MockDataCoord struct {
+	dcc.Base
+	initErr  error
+	startErr error
+	stopErr  error
+	regErr   error
+	stateErr commonpb.ErrorCode
+}
+
+func (m *MockDataCoord) Init() error {
+	return m.initErr
+}
+
+func (m *MockDataCoord) Start() error {
+	return m.startErr
+}
+
+func (m *MockDataCoord) Stop() error {
+	return m.stopErr
+}
+
+func (m *MockDataCoord) Register() error {
+	return m.regErr
+}
+
+func (m *MockDataCoord) GetComponentStates(ctx context.Context) (*internalpb.ComponentStates, error) {
+	return &internalpb.ComponentStates{
+		State:  &internalpb.ComponentInfo{StateCode: internalpb.StateCode_Healthy},
+		Status: &commonpb.Status{ErrorCode: m.stateErr},
+	}, nil
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
+func Test_NewServer(t *testing.T) {
+	ctx := context.Background()
+	qcs, err := NewServer(ctx, nil)
+	assert.Nil(t, err)
+	assert.NotNil(t, qcs)
+
+	mqc := &MockQueryCoord{
+		states: &internalpb.ComponentStates{
+			State:  &internalpb.ComponentInfo{StateCode: internalpb.StateCode_Healthy},
+			Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success},
+		},
+		status:       &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success},
+		err:          nil,
+		strResp:      &milvuspb.StringResponse{Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}},
+		showcolResp:  &querypb.ShowCollectionsResponse{},
+		showpartResp: &querypb.ShowPartitionsResponse{},
+		partResp:     &querypb.GetPartitionStatesResponse{},
+		channelResp:  &querypb.CreateQueryChannelResponse{},
+		infoResp:     &querypb.GetSegmentInfoResponse{Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}},
+		metricResp:   &milvuspb.GetMetricsResponse{Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}},
+	}
+
+	mdc := &MockDataCoord{
+		stateErr: commonpb.ErrorCode_Success,
+	}
+
+	mrc := &MockRootCoord{
+		stateErr: commonpb.ErrorCode_Success,
+	}
+
+	t.Run("Run", func(t *testing.T) {
+		qcs.queryCoord = mqc
+		qcs.dataCoord = mdc
+		qcs.rootCoord = mrc
+
+		err = qcs.Run()
+		assert.Nil(t, err)
+	})
+
+	t.Run("GetComponentStates", func(t *testing.T) {
+		req := &internalpb.GetComponentStatesRequest{}
+		states, err := qcs.GetComponentStates(ctx, req)
+		assert.Nil(t, err)
+		assert.Equal(t, internalpb.StateCode_Healthy, states.State.StateCode)
+	})
+
+	t.Run("GetStatisticsChannel", func(t *testing.T) {
+		req := &internalpb.GetStatisticsChannelRequest{}
+		resp, err := qcs.GetStatisticsChannel(ctx, req)
+		assert.Nil(t, err)
+		assert.Equal(t, commonpb.ErrorCode_Success, resp.Status.ErrorCode)
+	})
+
+	t.Run("GetTimeTickChannel", func(t *testing.T) {
+		req := &internalpb.GetTimeTickChannelRequest{}
+		resp, err := qcs.GetTimeTickChannel(ctx, req)
+		assert.Nil(t, err)
+		assert.Equal(t, commonpb.ErrorCode_Success, resp.Status.ErrorCode)
+	})
+
+	t.Run("ShowCollections", func(t *testing.T) {
+		resp, err := qcs.ShowCollections(ctx, nil)
+		assert.Nil(t, err)
+		assert.NotNil(t, resp)
+	})
+
+	t.Run("LoadCollection", func(t *testing.T) {
+		resp, err := qcs.LoadCollection(ctx, nil)
+		assert.Nil(t, err)
+		assert.NotNil(t, resp)
+	})
+
+	t.Run("ReleaseCollection", func(t *testing.T) {
+		resp, err := qcs.ReleaseCollection(ctx, nil)
+		assert.Nil(t, err)
+		assert.NotNil(t, resp)
+	})
+
+	t.Run("ShowPartitions", func(t *testing.T) {
+		resp, err := qcs.ShowPartitions(ctx, nil)
+		assert.Nil(t, err)
+		assert.NotNil(t, resp)
+	})
+	t.Run("GetPartitionStates", func(t *testing.T) {
+		resp, err := qcs.GetPartitionStates(ctx, nil)
+		assert.Nil(t, err)
+		assert.NotNil(t, resp)
+	})
+
+	t.Run("LoadPartitions", func(t *testing.T) {
+		resp, err := qcs.LoadPartitions(ctx, nil)
+		assert.Nil(t, err)
+		assert.NotNil(t, resp)
+	})
+
+	t.Run("ReleasePartitions", func(t *testing.T) {
+		resp, err := qcs.ReleasePartitions(ctx, nil)
+		assert.Nil(t, err)
+		assert.NotNil(t, resp)
+	})
+
+	t.Run("CreateQueryChannel", func(t *testing.T) {
+		resp, err := qcs.CreateQueryChannel(ctx, nil)
+		assert.Nil(t, err)
+		assert.NotNil(t, resp)
+	})
+
+	t.Run("GetTimeTickChannel", func(t *testing.T) {
+		resp, err := qcs.GetTimeTickChannel(ctx, nil)
+		assert.Nil(t, err)
+		assert.NotNil(t, resp)
+	})
+
+	t.Run("GetSegmentInfo", func(t *testing.T) {
+		req := &querypb.GetSegmentInfoRequest{}
+		resp, err := qcs.GetSegmentInfo(ctx, req)
+		assert.Nil(t, err)
+		assert.Equal(t, commonpb.ErrorCode_Success, resp.Status.ErrorCode)
+	})
+
+	t.Run("GetMetrics", func(t *testing.T) {
+		req := &milvuspb.GetMetricsRequest{
+			Request: "",
+		}
+		resp, err := qcs.GetMetrics(ctx, req)
+		assert.Nil(t, err)
+		assert.Equal(t, commonpb.ErrorCode_Success, resp.Status.ErrorCode)
+	})
+
+	err = qcs.Stop()
+	assert.Nil(t, err)
+}
+
+func TestServer_Run1(t *testing.T) {
+	ctx := context.Background()
+	qcs, err := NewServer(ctx, nil)
+	assert.Nil(t, err)
+	assert.NotNil(t, qcs)
+
+	qcs.queryCoord = &MockQueryCoord{
+		regErr: errors.New("error"),
+	}
+	err = qcs.Run()
+	assert.Error(t, err)
+
+	err = qcs.Stop()
+	assert.Nil(t, err)
+}
+
+func TestServer_Run2(t *testing.T) {
+	ctx := context.Background()
+	qcs, err := NewServer(ctx, nil)
+	assert.Nil(t, err)
+	assert.NotNil(t, qcs)
+
+	qcs.queryCoord = &MockQueryCoord{}
+	qcs.rootCoord = &MockRootCoord{
+		initErr: errors.New("error"),
+	}
+	assert.Panics(t, func() { qcs.Run() })
+	err = qcs.Stop()
+	assert.Nil(t, err)
+}
+
+func TestServer_Run3(t *testing.T) {
+	ctx := context.Background()
+	qcs, err := NewServer(ctx, nil)
+	assert.Nil(t, err)
+	assert.NotNil(t, qcs)
+
+	qcs.queryCoord = &MockQueryCoord{}
+	qcs.rootCoord = &MockRootCoord{
+		startErr: errors.New("error"),
+	}
+	assert.Panics(t, func() { qcs.Run() })
+	err = qcs.Stop()
+	assert.Nil(t, err)
+
+}
+
+func TestServer_Run4(t *testing.T) {
+	ctx := context.Background()
+	qcs, err := NewServer(ctx, nil)
+	assert.Nil(t, err)
+	assert.NotNil(t, qcs)
+
+	qcs.queryCoord = &MockQueryCoord{}
+	qcs.rootCoord = &MockRootCoord{}
+	qcs.dataCoord = &MockDataCoord{
+		initErr: errors.New("error"),
+	}
+	assert.Panics(t, func() { qcs.Run() })
+	err = qcs.Stop()
+	assert.Nil(t, err)
+}
+
+func TestServer_Run5(t *testing.T) {
+	ctx := context.Background()
+	qcs, err := NewServer(ctx, nil)
+	assert.Nil(t, err)
+	assert.NotNil(t, qcs)
+
+	qcs.queryCoord = &MockQueryCoord{}
+	qcs.rootCoord = &MockRootCoord{}
+	qcs.dataCoord = &MockDataCoord{
+		startErr: errors.New("error"),
+	}
+	assert.Panics(t, func() { qcs.Run() })
+	err = qcs.Stop()
+	assert.Nil(t, err)
+}
--- a/internal/types/types.go
+++ b/internal/types/types.go
@ -550,3 +550,18 @@ type QueryCoord interface {

 	GetMetrics(ctx context.Context, req *milvuspb.GetMetricsRequest) (*milvuspb.GetMetricsResponse, error)
 }
+
+// QueryCoordComponent is used by grpc server of QueryCoord
+type QueryCoordComponent interface {
+	QueryCoord
+
+	// UpdateStateCode updates state code for QueryCoord
+	// State includes: Initializing, Healthy and Abnormal
+	UpdateStateCode(internalpb.StateCode)
+
+	// SetDataCoord set DataCoord for QueryCoord
+	SetDataCoord(DataCoord)
+
+	// SetRootCoord set RootCoord for QueryCoord
+	SetRootCoord(RootCoord)
+}