Fix datacoord_test.go failed (#6558)

Create a mock rootcoord when testing datacoord to hide the network layer
issue: #6557
Signed-off-by: sunby <bingyi.sun@zilliz.com>

internal/datacoord/server.go

@@ -90,8 +90,16 @@ type Server struct {
 	rootCoordClientCreator rootCoordCreatorFunc
 }
 
+type Option func(svr *Server)
+
+func SetRootCoordCreator(creator rootCoordCreatorFunc) Option {
+	return func(svr *Server) {
+		svr.rootCoordClientCreator = creator
+	}
+}
+
 // CreateServer create `Server` instance
-func CreateServer(ctx context.Context, factory msgstream.Factory) (*Server, error) {
+func CreateServer(ctx context.Context, factory msgstream.Factory, opts ...Option) (*Server, error) {
 	rand.Seed(time.Now().UnixNano())
 	s := &Server{
 		ctx:                    ctx,
@@ -100,6 +108,10 @@ func CreateServer(ctx context.Context, factory msgstream.Factory) (*Server, erro
 		dataClientCreator:      defaultDataNodeCreatorFunc,
 		rootCoordClientCreator: defaultRootCoordCreatorFunc,
 	}
+
+	for _, opt := range opts {
+		opt(s)
+	}
 	return s, nil
 }
 

internal/distributed/datacoord/datacoord_test.go

@@ -18,6 +18,7 @@ import (
 	"math/rand"
 	"testing"
 
+	"github.com/milvus-io/milvus/internal/datacoord"
 	"github.com/milvus-io/milvus/internal/msgstream"
 	"github.com/milvus-io/milvus/internal/proto/commonpb"
 	"github.com/milvus-io/milvus/internal/proto/datapb"
@@ -82,18 +83,21 @@ func (m *mockRootCoord) ShowCollections(ctx context.Context, req *milvuspb.ShowC
 func TestRun(t *testing.T) {
 	ctx := context.Background()
 	msFactory := msgstream.NewPmsFactory()
-	dsServer, err := NewServer(ctx, msFactory)
+	var mockRootCoordCreator = func(ctx context.Context, metaRootPath string,
+		etcdEndpoints []string) (types.RootCoord, error) {
+		return &mockRootCoord{}, nil
+	}
+	dsServer, err := NewServer(ctx, msFactory, datacoord.SetRootCoordCreator(mockRootCoordCreator))
 	assert.Nil(t, err)
 
 	Params.Init()
-	Params.Port = 1000000
-	err = dsServer.Run()
-	assert.NotNil(t, err)
-	assert.EqualError(t, err, "listen tcp: address 1000000: invalid port")
-
 	Params.Port = rand.Int()%100 + 10000
 	err = dsServer.Run()
 	assert.Nil(t, err)
+	defer func() {
+		err = dsServer.Stop()
+		assert.Nil(t, err)
+	}()
 
 	t.Run("get component states", func(t *testing.T) {
 		req := &internalpb.GetComponentStatesRequest{}
@@ -162,7 +166,7 @@ func TestRun(t *testing.T) {
 		req := &datapb.GetPartitionStatisticsRequest{}
 		rsp, err := dsServer.GetPartitionStatistics(ctx, req)
 		assert.Nil(t, err)
-		assert.Nil(t, rsp)
+		assert.Equal(t, rsp.Status.ErrorCode, commonpb.ErrorCode_Success)
 	})
 
 	t.Run("get segment info channel", func(t *testing.T) {
@@ -172,7 +176,5 @@ func TestRun(t *testing.T) {
 		assert.Equal(t, rsp.Status.ErrorCode, commonpb.ErrorCode_Success)
 	})
 
-	err = dsServer.Stop()
-	assert.Nil(t, err)
 }
 */

internal/distributed/datacoord/service.go

@@ -30,7 +30,6 @@ import (
 	"github.com/milvus-io/milvus/internal/datacoord"
 	"github.com/milvus-io/milvus/internal/log"
 	"github.com/milvus-io/milvus/internal/msgstream"
-	"github.com/milvus-io/milvus/internal/types"
 	"github.com/milvus-io/milvus/internal/util/funcutil"
 	"github.com/milvus-io/milvus/internal/util/trace"
 
@@ -41,21 +40,19 @@ import (
 )
 
 type Server struct {
+	ctx    context.Context
+	cancel context.CancelFunc
+
+	wg        sync.WaitGroup
 	dataCoord *datacoord.Server
-	ctx       context.Context
-	cancel    context.CancelFunc
 
 	grpcErrChan chan error
-	wg          sync.WaitGroup
-
-	grpcServer *grpc.Server
-	rootCoord  types.RootCoord
-
-	closer io.Closer
+	grpcServer  *grpc.Server
+	closer      io.Closer
 }
 
 // NewServer new data service grpc server
-func NewServer(ctx context.Context, factory msgstream.Factory) (*Server, error) {
+func NewServer(ctx context.Context, factory msgstream.Factory, opts ...datacoord.Option) (*Server, error) {
 	var err error
 	ctx1, cancel := context.WithCancel(ctx)
 
@@ -64,7 +61,7 @@ func NewServer(ctx context.Context, factory msgstream.Factory) (*Server, error)
 		cancel:      cancel,
 		grpcErrChan: make(chan error),
 	}
-	s.dataCoord, err = datacoord.CreateServer(s.ctx, factory)
+	s.dataCoord, err = datacoord.CreateServer(s.ctx, factory, opts...)
 	if err != nil {
 		return nil, err
 	}