Improve code coverage for distributed/datanode (#5133)

Signed-off-by: yudong.cai <yudong.cai@zilliz.com>

cmd/distributed/components/data_node.go

@@ -25,8 +25,7 @@ type DataNode struct {
 }
 
 func NewDataNode(ctx context.Context, factory msgstream.Factory) (*DataNode, error) {
-
-	svr, err := grpcdatanode.New(ctx, factory)
+	svr, err := grpcdatanode.NewServer(ctx, factory)
 	if err != nil {
 		return nil, err
 	}

internal/datanode/data_sync_service.go

@@ -58,7 +58,11 @@ func (dsService *dataSyncService) init() {
 
 func (dsService *dataSyncService) start() {
 	log.Debug("Data Sync Service Start Successfully")
-	dsService.fg.Start()
+	if dsService.fg != nil {
+		dsService.fg.Start()
+	} else {
+		log.Debug("Data Sync Service flowgraph nil")
+	}
 }
 
 func (dsService *dataSyncService) close() {

internal/distributed/datanode/datanode_test.go

@@ -0,0 +1,172 @@
+// 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 grpcdatanode
+
+import (
+	"context"
+	"fmt"
+	"math/rand"
+	"net"
+	"strconv"
+	"testing"
+
+	"github.com/milvus-io/milvus/internal/msgstream"
+	"github.com/milvus-io/milvus/internal/proto/commonpb"
+	"github.com/milvus-io/milvus/internal/proto/datapb"
+	"github.com/milvus-io/milvus/internal/proto/internalpb"
+	"github.com/milvus-io/milvus/internal/proto/milvuspb"
+	"github.com/milvus-io/milvus/internal/types"
+	"github.com/stretchr/testify/assert"
+)
+
+type mockMaster struct {
+	types.MasterService
+}
+
+func (m *mockMaster) Init() error {
+	return nil
+}
+
+func (m *mockMaster) Start() error {
+	return nil
+}
+
+func (m *mockMaster) Stop() error {
+	return fmt.Errorf("stop error")
+}
+
+func (m *mockMaster) GetComponentStates(ctx context.Context) (*internalpb.ComponentStates, error) {
+	return &internalpb.ComponentStates{
+		State: &internalpb.ComponentInfo{
+			StateCode: internalpb.StateCode_Healthy,
+		},
+		Status: &commonpb.Status{
+			ErrorCode: commonpb.ErrorCode_Success,
+		},
+		SubcomponentStates: []*internalpb.ComponentInfo{
+			{
+				StateCode: internalpb.StateCode_Healthy,
+			},
+		},
+	}, nil
+}
+
+func (m *mockMaster) ShowCollections(ctx context.Context, req *milvuspb.ShowCollectionsRequest) (*milvuspb.ShowCollectionsResponse, error) {
+	return &milvuspb.ShowCollectionsResponse{
+		Status: &commonpb.Status{
+			ErrorCode: commonpb.ErrorCode_Success,
+			Reason:    "",
+		},
+		CollectionNames: []string{},
+	}, nil
+}
+
+type mockDataService struct {
+	types.DataService
+}
+
+func (m *mockDataService) Init() error {
+	return nil
+}
+
+func (m *mockDataService) Start() error {
+	return nil
+}
+
+func (m *mockDataService) Stop() error {
+	return fmt.Errorf("stop error")
+}
+
+func (m *mockDataService) GetComponentStates(ctx context.Context) (*internalpb.ComponentStates, error) {
+	return &internalpb.ComponentStates{
+		State: &internalpb.ComponentInfo{
+			StateCode: internalpb.StateCode_Healthy,
+		},
+		Status: &commonpb.Status{
+			ErrorCode: commonpb.ErrorCode_Success,
+		},
+		SubcomponentStates: []*internalpb.ComponentInfo{
+			{
+				StateCode: internalpb.StateCode_Healthy,
+			},
+		},
+	}, nil
+}
+
+func (m *mockDataService) RegisterNode(ctx context.Context, req *datapb.RegisterNodeRequest) (*datapb.RegisterNodeResponse, error) {
+	return &datapb.RegisterNodeResponse{
+		Status: &commonpb.Status{
+			ErrorCode: commonpb.ErrorCode_Success,
+		},
+		InitParams: &internalpb.InitParams{
+			NodeID: int64(1),
+		},
+	}, nil
+}
+
+func TestRun(t *testing.T) {
+	ctx := context.Background()
+	msFactory := msgstream.NewPmsFactory()
+	dnServer, err := NewServer(ctx, msFactory)
+	assert.Nil(t, err)
+
+	Params.Init()
+
+	dnServer.newMasterServiceClient = func(s string) (types.MasterService, error) {
+		return &mockMaster{}, nil
+	}
+	dnServer.newDataServiceClient = func(s string) types.DataService {
+		return &mockDataService{}
+	}
+
+	grpcPort := rand.Int()%100 + 10000
+	Params.listener, _ = net.Listen("tcp", ":"+strconv.Itoa(grpcPort))
+
+	// to let datanode init pass
+	dnServer.datanode.UpdateStateCode(internalpb.StateCode_Initializing)
+	dnServer.datanode.WatchDmChannels(ctx, nil)
+
+	err = dnServer.Run()
+	assert.Nil(t, err)
+
+	t.Run("get component states", func(t *testing.T) {
+		req := &internalpb.GetComponentStatesRequest{}
+		rsp, err := dnServer.GetComponentStates(ctx, req)
+		assert.Nil(t, err)
+		assert.Equal(t, rsp.Status.ErrorCode, commonpb.ErrorCode_Success)
+	})
+
+	t.Run("get statistics channel", func(t *testing.T) {
+		req := &internalpb.GetStatisticsChannelRequest{}
+		rsp, err := dnServer.GetStatisticsChannel(ctx, req)
+		assert.Nil(t, err)
+		assert.Equal(t, rsp.Status.ErrorCode, commonpb.ErrorCode_Success)
+	})
+
+	t.Run("watch dm channels", func(t *testing.T) {
+		req := &datapb.WatchDmChannelsRequest{}
+		_, err := dnServer.WatchDmChannels(ctx, req)
+		assert.NotNil(t, err)
+	})
+
+	t.Run("flush segments", func(t *testing.T) {
+		req := &datapb.FlushSegmentsRequest{
+			Base: &commonpb.MsgBase{},
+		}
+		rsp, err := dnServer.FlushSegments(ctx, req)
+		assert.Nil(t, err)
+		assert.Equal(t, rsp.ErrorCode, commonpb.ErrorCode_Success)
+	})
+
+	err = dnServer.Stop()
+	assert.Nil(t, err)
+}

internal/distributed/datanode/service.go

@@ -55,16 +55,26 @@ type Server struct {
 	masterService types.MasterService
 	dataService   types.DataService
 
+	newMasterServiceClient func(string) (types.MasterService, error)
+	newDataServiceClient   func(string) types.DataService
+
 	closer io.Closer
 }
 
-func New(ctx context.Context, factory msgstream.Factory) (*Server, error) {
+// NewServer new data node grpc server
+func NewServer(ctx context.Context, factory msgstream.Factory) (*Server, error) {
 	ctx1, cancel := context.WithCancel(ctx)
 	var s = &Server{
 		ctx:         ctx1,
 		cancel:      cancel,
 		msFactory:   factory,
 		grpcErrChan: make(chan error),
+		newMasterServiceClient: func(s string) (types.MasterService, error) {
+			return msc.NewClient(s, 20*time.Second)
+		},
+		newDataServiceClient: func(s string) types.DataService {
+			return dsc.NewClient(Params.DataServiceAddress)
+		},
 	}
 
 	s.datanode = dn.NewDataNode(s.ctx, s.msFactory)
@@ -72,6 +82,14 @@ func New(ctx context.Context, factory msgstream.Factory) (*Server, error) {
 	return s, nil
 }
 
+func (s *Server) startGrpc() error {
+	s.wg.Add(1)
+	go s.startGrpcLoop(Params.listener)
+	// wait for grpc server loop start
+	err := <-s.grpcErrChan
+	return err
+}
+
 func (s *Server) startGrpcLoop(listener net.Listener) {
 	defer s.wg.Done()
 
@@ -105,7 +123,6 @@ func (s *Server) SetDataServiceInterface(ds types.DataService) error {
 }
 
 func (s *Server) Run() error {
-
 	if err := s.init(); err != nil {
 		return err
 	}
@@ -152,55 +169,52 @@ func (s *Server) init() error {
 	addr := Params.IP + ":" + strconv.Itoa(Params.Port)
 	log.Debug("DataNode address", zap.String("address", addr))
 
-	s.wg.Add(1)
-	go s.startGrpcLoop(Params.listener)
-	// wait for grpc server loop start
-	err := <-s.grpcErrChan
+	err := s.startGrpc()
 	if err != nil {
 		return err
 	}
 
 	// --- Master Server Client ---
-	log.Debug("Master service address", zap.String("address", Params.MasterAddress))
-	log.Debug("Init master service client ...")
-	masterClient, err := msc.NewClient(Params.MasterAddress, 20*time.Second)
-	if err != nil {
-		panic(err)
-	}
-
-	if err = masterClient.Init(); err != nil {
-		panic(err)
-	}
-
-	if err = masterClient.Start(); err != nil {
-		panic(err)
-	}
-	err = funcutil.WaitForComponentHealthy(ctx, masterClient, "MasterService", 1000000, time.Millisecond*200)
-
-	if err != nil {
-		panic(err)
-	}
-
-	if err := s.SetMasterServiceInterface(masterClient); err != nil {
-		panic(err)
+	if s.newMasterServiceClient != nil {
+		log.Debug("Master service address", zap.String("address", Params.MasterAddress))
+		log.Debug("Init master service client ...")
+		masterServiceClient, err := s.newMasterServiceClient(Params.MasterAddress)
+		if err != nil {
+			panic(err)
+		}
+		if err = masterServiceClient.Init(); err != nil {
+			panic(err)
+		}
+		if err = masterServiceClient.Start(); err != nil {
+			panic(err)
+		}
+		err = funcutil.WaitForComponentHealthy(ctx, masterServiceClient, "MasterService", 1000000, time.Millisecond*200)
+		if err != nil {
+			panic(err)
+		}
+		if err = s.SetMasterServiceInterface(masterServiceClient); err != nil {
+			panic(err)
+		}
 	}
 
 	// --- Data Server Client ---
-	log.Debug("Data service address", zap.String("address", Params.DataServiceAddress))
-	log.Debug("DataNode Init data service client ...")
-	dataService := dsc.NewClient(Params.DataServiceAddress)
-	if err = dataService.Init(); err != nil {
-		panic(err)
-	}
-	if err = dataService.Start(); err != nil {
-		panic(err)
-	}
-	err = funcutil.WaitForComponentInitOrHealthy(ctx, dataService, "DataService", 1000000, time.Millisecond*200)
-	if err != nil {
-		panic(err)
-	}
-	if err := s.SetDataServiceInterface(dataService); err != nil {
-		panic(err)
+	if s.newDataServiceClient != nil {
+		log.Debug("Data service address", zap.String("address", Params.DataServiceAddress))
+		log.Debug("DataNode Init data service client ...")
+		dataServiceClient := s.newDataServiceClient(Params.DataServiceAddress)
+		if err = dataServiceClient.Init(); err != nil {
+			panic(err)
+		}
+		if err = dataServiceClient.Start(); err != nil {
+			panic(err)
+		}
+		err = funcutil.WaitForComponentInitOrHealthy(ctx, dataServiceClient, "DataService", 1000000, time.Millisecond*200)
+		if err != nil {
+			panic(err)
+		}
+		if err = s.SetDataServiceInterface(dataServiceClient); err != nil {
+			panic(err)
+		}
 	}
 
 	s.datanode.NodeID = dn.Params.NodeID

internal/distributed/dataservice/dataservice_test.go

@@ -19,6 +19,7 @@ import (
 
 	"github.com/milvus-io/milvus/internal/msgstream"
 	"github.com/milvus-io/milvus/internal/proto/commonpb"
+	"github.com/milvus-io/milvus/internal/proto/datapb"
 	"github.com/milvus-io/milvus/internal/proto/internalpb"
 	"github.com/milvus-io/milvus/internal/proto/milvuspb"
 	"github.com/milvus-io/milvus/internal/types"
@@ -80,23 +81,124 @@ func (m *mockMaster) ShowCollections(ctx context.Context, req *milvuspb.ShowColl
 func TestRun(t *testing.T) {
 	ctx := context.Background()
 	msFactory := msgstream.NewPmsFactory()
-	svr, err := NewServer(ctx, msFactory)
+	dsServer, err := NewServer(ctx, msFactory)
 	assert.Nil(t, err)
 
 	Params.Init()
 	Params.Port = 1000000
-	err = svr.Run()
+	err = dsServer.Run()
 	assert.NotNil(t, err)
 	assert.EqualError(t, err, "listen tcp: address 1000000: invalid port")
 
-	svr.newMasterServiceClient = func(s string) (types.MasterService, error) {
+	dsServer.newMasterServiceClient = func(s string) (types.MasterService, error) {
 		return &mockMaster{}, nil
 	}
 
 	Params.Port = rand.Int()%100 + 10000
-	err = svr.Run()
+	err = dsServer.Run()
 	assert.Nil(t, err)
 
-	err = svr.Stop()
+	t.Run("get component states", func(t *testing.T) {
+		req := &internalpb.GetComponentStatesRequest{}
+		rsp, err := dsServer.GetComponentStates(ctx, req)
+		assert.Nil(t, err)
+		assert.Equal(t, rsp.Status.ErrorCode, commonpb.ErrorCode_Success)
+	})
+
+	t.Run("get time tick channel", func(t *testing.T) {
+		req := &internalpb.GetTimeTickChannelRequest{}
+		rsp, err := dsServer.GetTimeTickChannel(ctx, req)
+		assert.Nil(t, err)
+		assert.Equal(t, rsp.Status.ErrorCode, commonpb.ErrorCode_Success)
+	})
+
+	t.Run("get statistics channel", func(t *testing.T) {
+		req := &internalpb.GetStatisticsChannelRequest{}
+		rsp, err := dsServer.GetStatisticsChannel(ctx, req)
+		assert.Nil(t, err)
+		assert.Equal(t, rsp.Status.ErrorCode, commonpb.ErrorCode_Success)
+	})
+
+	t.Run("get segment info", func(t *testing.T) {
+		req := &datapb.GetSegmentInfoRequest{}
+		rsp, err := dsServer.GetSegmentInfo(ctx, req)
+		assert.Nil(t, err)
+		assert.Equal(t, rsp.Status.ErrorCode, commonpb.ErrorCode_Success)
+	})
+
+	//t.Run("register node", func(t *testing.T) {
+	//	req := &datapb.RegisterNodeRequest{
+	//		Base: &commonpb.MsgBase{},
+	//		Address: &commonpb.Address{},
+	//	}
+	//	rsp, err := dsServer.RegisterNode(ctx, req)
+	//	assert.Nil(t, err)
+	//	assert.Equal(t, rsp.Status.ErrorCode, commonpb.ErrorCode_Success)
+	//})
+
+	t.Run("flush", func(t *testing.T) {
+		req := &datapb.FlushRequest{}
+		rsp, err := dsServer.Flush(ctx, req)
+		assert.Nil(t, err)
+		assert.Equal(t, rsp.ErrorCode, commonpb.ErrorCode_Success)
+	})
+
+	t.Run("assign segment id", func(t *testing.T) {
+		req := &datapb.AssignSegmentIDRequest{}
+		rsp, err := dsServer.AssignSegmentID(ctx, req)
+		assert.Nil(t, err)
+		assert.Equal(t, rsp.Status.ErrorCode, commonpb.ErrorCode_Success)
+	})
+
+	t.Run("show segments", func(t *testing.T) {
+		req := &datapb.ShowSegmentsRequest{}
+		rsp, err := dsServer.ShowSegments(ctx, req)
+		assert.Nil(t, err)
+		assert.Equal(t, rsp.Status.ErrorCode, commonpb.ErrorCode_Success)
+	})
+
+	t.Run("get segment states", func(t *testing.T) {
+		req := &datapb.GetSegmentStatesRequest{}
+		rsp, err := dsServer.GetSegmentStates(ctx, req)
+		assert.Nil(t, err)
+		assert.Equal(t, rsp.Status.ErrorCode, commonpb.ErrorCode_Success)
+	})
+
+	t.Run("get insert binlog paths", func(t *testing.T) {
+		req := &datapb.GetInsertBinlogPathsRequest{}
+		rsp, err := dsServer.GetInsertBinlogPaths(ctx, req)
+		assert.Nil(t, err)
+		assert.Equal(t, rsp.Status.ErrorCode, commonpb.ErrorCode_Success)
+	})
+
+	t.Run("get insert channels", func(t *testing.T) {
+		req := &datapb.GetInsertChannelsRequest{}
+		rsp, err := dsServer.GetInsertChannels(ctx, req)
+		assert.Nil(t, err)
+		assert.Equal(t, rsp.Status.ErrorCode, commonpb.ErrorCode_Success)
+	})
+
+	t.Run("get collection statistics", func(t *testing.T) {
+		req := &datapb.GetCollectionStatisticsRequest{}
+		rsp, err := dsServer.GetCollectionStatistics(ctx, req)
+		assert.Nil(t, err)
+		assert.Equal(t, rsp.Status.ErrorCode, commonpb.ErrorCode_Success)
+	})
+
+	t.Run("get partition statistics", func(t *testing.T) {
+		req := &datapb.GetPartitionStatisticsRequest{}
+		rsp, err := dsServer.GetPartitionStatistics(ctx, req)
+		assert.Nil(t, err)
+		assert.Nil(t, rsp)
+	})
+
+	t.Run("get segment info channel", func(t *testing.T) {
+		req := &datapb.GetSegmentInfoChannelRequest{}
+		rsp, err := dsServer.GetSegmentInfoChannel(ctx, req)
+		assert.Nil(t, err)
+		assert.Equal(t, rsp.Status.ErrorCode, commonpb.ErrorCode_Success)
+	})
+
+	err = dsServer.Stop()
 	assert.Nil(t, err)
 }

internal/distributed/dataservice/service.go

@@ -58,6 +58,7 @@ type Server struct {
 	closer io.Closer
 }
 
+// NewServer new data service grpc server
 func NewServer(ctx context.Context, factory msgstream.Factory) (*Server, error) {
 	var err error
 	ctx1, cancel := context.WithCancel(ctx)
@@ -188,7 +189,6 @@ func (s *Server) Stop() error {
 }
 
 func (s *Server) Run() error {
-
 	if err := s.init(); err != nil {
 		return err
 	}