fix: ut

Signed-off-by: chyezh <[email protected]>
milvus-io · Jan 7, 2025 · 9ff7e5b · 9ff7e5b
1 parent 4104b3c
commit 9ff7e5b
Show file tree

Hide file tree

Showing 7 changed files with 60 additions and 39 deletions.
diff --git a/internal/coordinator/coordclient/registry.go b/internal/coordinator/coordclient/registry.go
@@ -58,20 +58,29 @@ func EnableLocalClientRole(cfg *LocalClientRoleConfig) {
 
 // RegisterQueryCoordServer register query coord server
 func RegisterQueryCoordServer(server querypb.QueryCoordServer) {
+	if !enableLocal.EnableQueryCoord {
+		return
+	}
 	newLocalClient := grpcclient.NewLocalGRPCClient(&querypb.QueryCoord_ServiceDesc, server, querypb.NewQueryCoordClient)
 	glocalClient.queryCoordClient.Set(&nopCloseQueryCoordClient{newLocalClient})
 	log.Ctx(context.TODO()).Info("register query coord server", zap.Any("enableLocalClient", enableLocal))
 }
 
 // RegsterDataCoordServer register data coord server
 func RegisterDataCoordServer(server datapb.DataCoordServer) {
+	if !enableLocal.EnableDataCoord {
+		return
+	}
 	newLocalClient := grpcclient.NewLocalGRPCClient(&datapb.DataCoord_ServiceDesc, server, datapb.NewDataCoordClient)
 	glocalClient.dataCoordClient.Set(&nopCloseDataCoordClient{newLocalClient})
 	log.Ctx(context.TODO()).Info("register data coord server", zap.Any("enableLocalClient", enableLocal))
 }
 
 // RegisterRootCoordServer register root coord server
 func RegisterRootCoordServer(server rootcoordpb.RootCoordServer) {
+	if !enableLocal.EnableRootCoord {
+		return
+	}
 	newLocalClient := grpcclient.NewLocalGRPCClient(&rootcoordpb.RootCoord_ServiceDesc, server, rootcoordpb.NewRootCoordClient)
 	glocalClient.rootCoordClient.Set(&nopCloseRootCoordClient{newLocalClient})
 	log.Ctx(context.TODO()).Info("register root coord server", zap.Any("enableLocalClient", enableLocal))
@@ -128,9 +137,6 @@ func GetRootCoordClient(ctx context.Context) types.RootCoordClient {
 // MustGetLocalRootCoordClientFuture return root coord client future,
 // panic if root coord client is not enabled
 func MustGetLocalRootCoordClientFuture() *syncutil.Future[types.RootCoordClient] {
-	if !enableLocal.EnableRootCoord {
-		panic("root coord client is not enabled")
-	}
 	return glocalClient.rootCoordClient
 }
 

diff --git a/internal/coordinator/coordclient/registry_test.go b/internal/coordinator/coordclient/registry_test.go
@@ -31,6 +31,13 @@ func TestRegistry(t *testing.T) {
 	assert.False(t, enableLocal.EnableDataCoord)
 	assert.False(t, enableLocal.EnableRootCoord)
 
+	RegisterRootCoordServer(&rootcoordpb.UnimplementedRootCoordServer{})
+	RegisterDataCoordServer(&datapb.UnimplementedDataCoordServer{})
+	RegisterQueryCoordServer(&querypb.UnimplementedQueryCoordServer{})
+	assert.False(t, glocalClient.dataCoordClient.Ready())
+	assert.False(t, glocalClient.queryCoordClient.Ready())
+	assert.False(t, glocalClient.rootCoordClient.Ready())
+
 	enableLocal = &LocalClientRoleConfig{}
 
 	EnableLocalClientRole(&LocalClientRoleConfig{

diff --git a/internal/distributed/rootcoord/service_test.go b/internal/distributed/rootcoord/service_test.go
@@ -27,6 +27,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/tikv/client-go/v2/txnkv"
 	clientv3 "go.etcd.io/etcd/client/v3"
+	"google.golang.org/grpc"
 
 	"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
 	"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
@@ -97,6 +98,9 @@ func (m *mockCore) SetQueryCoordClient(client types.QueryCoordClient) error {
 func (m *mockCore) SetProxyCreator(func(ctx context.Context, addr string, nodeID int64) (types.ProxyClient, error)) {
 }
 
+func (m *mockCore) RegisterStreamingCoordGRPCService(server *grpc.Server) {
+}
+
 func (m *mockCore) Register() error {
 	return nil
 }

diff --git a/internal/mocks/mock_datacoord.go b/internal/mocks/mock_datacoord.go
diff --git a/internal/mocks/mock_rootcoord.go b/internal/mocks/mock_rootcoord.go
diff --git a/internal/rootcoord/root_coord.go b/internal/rootcoord/root_coord.go
@@ -826,7 +826,9 @@ func (c *Core) Stop() error {
 	c.stopExecutor()
 	c.stopScheduler()
 
-	c.streamingCoord.Stop()
+	if c.streamingCoord != nil {
+		c.streamingCoord.Stop()
+	}
 	if c.proxyWatcher != nil {
 		c.proxyWatcher.Stop()
 	}

diff --git a/internal/streamingcoord/server/service/discover/discover_server_test.go b/internal/streamingcoord/server/service/discover/discover_server_test.go
@@ -8,13 +8,15 @@ import (
 	"github.com/stretchr/testify/mock"
 
 	"github.com/milvus-io/milvus/internal/mocks/streamingcoord/server/mock_balancer"
+	"github.com/milvus-io/milvus/internal/streamingcoord/server/resource"
 	"github.com/milvus-io/milvus/pkg/mocks/streaming/proto/mock_streamingpb"
 	"github.com/milvus-io/milvus/pkg/streaming/proto/streamingpb"
 	"github.com/milvus-io/milvus/pkg/streaming/util/types"
 	"github.com/milvus-io/milvus/pkg/util/typeutil"
 )
 
 func TestAssignmentDiscover(t *testing.T) {
+	resource.InitForTest()
 	b := mock_balancer.NewMockBalancer(t)
 	b.EXPECT().WatchChannelAssignments(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, cb func(typeutil.VersionInt64Pair, []types.PChannelInfoAssigned) error) error {
 		versions := []typeutil.VersionInt64Pair{