Add ut for querycoord's impl (#7728)

Signed-off-by: xige-16 <xi.ge@zilliz.com>

internal/querycoord/impl.go

@@ -368,7 +368,7 @@ func (qc *QueryCoord) ReleasePartitions(ctx context.Context, req *querypb.Releas
 
 	hasCollection := qc.meta.hasCollection(collectionID)
 	if !hasCollection {
-		log.Warn("release partitions end, query coordinator don't have the log of", zap.String("collectionID", fmt.Sprintln(collectionID)))
+		log.Warn("release partitions end, query coordinator don't have the log of", zap.Int64("collectionID", collectionID))
 		return status, nil
 	}
 
@@ -380,6 +380,19 @@ func (qc *QueryCoord) ReleasePartitions(ctx context.Context, req *querypb.Releas
 		return status, err
 	}
 
+	toReleasedPartitions := make([]UniqueID, 0)
+	for _, id := range partitionIDs {
+		hasPartition := qc.meta.hasPartition(collectionID, id)
+		if hasPartition {
+			toReleasedPartitions = append(toReleasedPartitions, id)
+		}
+	}
+	if len(toReleasedPartitions) == 0 {
+		log.Warn("release partitions end, query coordinator don't have the log of", zap.Int64s("partitionIDs", partitionIDs))
+		return status, nil
+	}
+
+	req.PartitionIDs = toReleasedPartitions
 	releasePartitionTask := &ReleasePartitionTask{
 		BaseTask: BaseTask{
 			ctx:              qc.loopCtx,
@@ -541,7 +554,7 @@ func (qc *QueryCoord) GetMetrics(ctx context.Context, req *milvuspb.GetMetricsRe
 				Reason:    err.Error(),
 			},
 			Response: "",
-		}, nil
+		}, err
 	}
 
 	log.Debug("QueryCoord.GetMetrics",
@@ -568,16 +581,18 @@ func (qc *QueryCoord) GetMetrics(ctx context.Context, req *milvuspb.GetMetricsRe
 
 		return metrics, err
 	}
-	log.Debug("QueryCoord.GetMetrics failed, request metric type is not implemented yet",
+	err = errors.New(metricsinfo.MsgUnimplementedMetric)
+	log.Debug("QueryCoord.GetMetrics failed",
 		zap.Int64("node_id", Params.QueryCoordID),
 		zap.String("req", req.Request),
-		zap.String("metric_type", metricType))
+		zap.String("metric_type", metricType),
+		zap.Error(err))
 
 	return &milvuspb.GetMetricsResponse{
 		Status: &commonpb.Status{
 			ErrorCode: commonpb.ErrorCode_UnexpectedError,
-			Reason:    metricsinfo.MsgUnimplementedMetric,
+			Reason:    err.Error(),
 		},
 		Response: "",
-	}, nil
+	}, err
 }

internal/querycoord/impl_test.go

@@ -23,6 +23,30 @@ func TestGrpcTask(t *testing.T) {
 	node, err := startQueryNodeServer(ctx)
 	assert.Nil(t, err)
 
+	t.Run("Test ShowParsOnNotLoadedCol", func(t *testing.T) {
+		res, err := queryCoord.ShowPartitions(ctx, &querypb.ShowPartitionsRequest{
+			Base: &commonpb.MsgBase{
+				MsgType: commonpb.MsgType_ShowCollections,
+			},
+			CollectionID: defaultCollectionID,
+			PartitionIDs: []UniqueID{defaultPartitionID},
+		})
+		assert.Equal(t, commonpb.ErrorCode_UnexpectedError, res.Status.ErrorCode)
+		assert.NotNil(t, err)
+	})
+
+	t.Run("Test LoadEmptyPartition", func(t *testing.T) {
+		status, err := queryCoord.LoadPartitions(ctx, &querypb.LoadPartitionsRequest{
+			Base: &commonpb.MsgBase{
+				MsgType: commonpb.MsgType_LoadPartitions,
+			},
+			CollectionID: defaultCollectionID,
+			Schema:       genCollectionSchema(defaultCollectionID, false),
+		})
+		assert.Equal(t, commonpb.ErrorCode_UnexpectedError, status.ErrorCode)
+		assert.NotNil(t, err)
+	})
+
 	t.Run("Test LoadPartition", func(t *testing.T) {
 		status, err := queryCoord.LoadPartitions(ctx, &querypb.LoadPartitionsRequest{
 			Base: &commonpb.MsgBase{
@@ -32,7 +56,7 @@ func TestGrpcTask(t *testing.T) {
 			PartitionIDs: []UniqueID{defaultPartitionID},
 			Schema:       genCollectionSchema(defaultCollectionID, false),
 		})
-		assert.Equal(t, status.ErrorCode, commonpb.ErrorCode_Success)
+		assert.Equal(t, commonpb.ErrorCode_Success, status.ErrorCode)
 		assert.Nil(t, err)
 	})
 
@@ -44,10 +68,22 @@ func TestGrpcTask(t *testing.T) {
 			CollectionID: defaultCollectionID,
 			PartitionIDs: []UniqueID{defaultPartitionID},
 		})
-		assert.Equal(t, res.Status.ErrorCode, commonpb.ErrorCode_Success)
+		assert.Equal(t, commonpb.ErrorCode_Success, res.Status.ErrorCode)
 		assert.Nil(t, err)
 	})
 
+	t.Run("Test ShowNotLoadedPartitions", func(t *testing.T) {
+		res, err := queryCoord.ShowPartitions(ctx, &querypb.ShowPartitionsRequest{
+			Base: &commonpb.MsgBase{
+				MsgType: commonpb.MsgType_ShowCollections,
+			},
+			CollectionID: defaultCollectionID,
+			PartitionIDs: []UniqueID{-1},
+		})
+		assert.Equal(t, commonpb.ErrorCode_UnexpectedError, res.Status.ErrorCode)
+		assert.NotNil(t, err)
+	})
+
 	t.Run("Test ShowAllPartitions", func(t *testing.T) {
 		res, err := queryCoord.ShowPartitions(ctx, &querypb.ShowPartitionsRequest{
 			Base: &commonpb.MsgBase{
@@ -55,7 +91,7 @@ func TestGrpcTask(t *testing.T) {
 			},
 			CollectionID: defaultCollectionID,
 		})
-		assert.Equal(t, res.Status.ErrorCode, commonpb.ErrorCode_Success)
+		assert.Equal(t, commonpb.ErrorCode_Success, res.Status.ErrorCode)
 		assert.Nil(t, err)
 	})
 
@@ -67,7 +103,7 @@ func TestGrpcTask(t *testing.T) {
 			CollectionID: defaultCollectionID,
 			PartitionIDs: []UniqueID{defaultPartitionID},
 		})
-		assert.Equal(t, res.Status.ErrorCode, commonpb.ErrorCode_Success)
+		assert.Equal(t, commonpb.ErrorCode_Success, res.Status.ErrorCode)
 		assert.Nil(t, err)
 	})
 
@@ -79,7 +115,20 @@ func TestGrpcTask(t *testing.T) {
 			CollectionID: defaultCollectionID,
 			Schema:       genCollectionSchema(defaultCollectionID, false),
 		})
-		assert.Equal(t, status.ErrorCode, commonpb.ErrorCode_Success)
+		assert.Equal(t, commonpb.ErrorCode_Success, status.ErrorCode)
+		assert.Nil(t, err)
+	})
+
+	t.Run("Test LoadParAfterLoadCol", func(t *testing.T) {
+		status, err := queryCoord.LoadPartitions(ctx, &querypb.LoadPartitionsRequest{
+			Base: &commonpb.MsgBase{
+				MsgType: commonpb.MsgType_LoadPartitions,
+			},
+			CollectionID: defaultCollectionID,
+			PartitionIDs: []UniqueID{defaultPartitionID},
+			Schema:       genCollectionSchema(defaultCollectionID, false),
+		})
+		assert.Equal(t, commonpb.ErrorCode_Success, status.ErrorCode)
 		assert.Nil(t, err)
 	})
 
@@ -90,18 +139,28 @@ func TestGrpcTask(t *testing.T) {
 			},
 			CollectionIDs: []UniqueID{defaultCollectionID},
 		})
-		assert.Equal(t, 100, int(res.InMemoryPercentages[0]))
-		assert.Equal(t, res.Status.ErrorCode, commonpb.ErrorCode_Success)
+		assert.Equal(t, commonpb.ErrorCode_Success, res.Status.ErrorCode)
 		assert.Nil(t, err)
 	})
 
+	t.Run("Test ShowNotLoadedCollections", func(t *testing.T) {
+		res, err := queryCoord.ShowCollections(ctx, &querypb.ShowCollectionsRequest{
+			Base: &commonpb.MsgBase{
+				MsgType: commonpb.MsgType_ShowCollections,
+			},
+			CollectionIDs: []UniqueID{-1},
+		})
+		assert.Equal(t, commonpb.ErrorCode_UnexpectedError, res.Status.ErrorCode)
+		assert.NotNil(t, err)
+	})
+
 	t.Run("Test ShowAllCollections", func(t *testing.T) {
 		res, err := queryCoord.ShowCollections(ctx, &querypb.ShowCollectionsRequest{
 			Base: &commonpb.MsgBase{
 				MsgType: commonpb.MsgType_ShowCollections,
 			},
 		})
-		assert.Equal(t, res.Status.ErrorCode, commonpb.ErrorCode_Success)
+		assert.Equal(t, commonpb.ErrorCode_Success, res.Status.ErrorCode)
 		assert.Nil(t, err)
 	})
 
@@ -112,7 +171,42 @@ func TestGrpcTask(t *testing.T) {
 			},
 			SegmentIDs: []UniqueID{defaultSegmentID},
 		})
-		assert.Equal(t, res.Status.ErrorCode, commonpb.ErrorCode_Success)
+		assert.Equal(t, commonpb.ErrorCode_Success, res.Status.ErrorCode)
+		assert.Nil(t, err)
+	})
+
+	t.Run("Test ReleaseParOfNotLoadedCol", func(t *testing.T) {
+		status, err := queryCoord.ReleasePartitions(ctx, &querypb.ReleasePartitionsRequest{
+			Base: &commonpb.MsgBase{
+				MsgType: commonpb.MsgType_ReleasePartitions,
+			},
+			CollectionID: -1,
+			PartitionIDs: []UniqueID{defaultPartitionID},
+		})
+		assert.Equal(t, commonpb.ErrorCode_Success, status.ErrorCode)
+		assert.Nil(t, err)
+	})
+
+	t.Run("Test ReleaseEmptyPartitions", func(t *testing.T) {
+		status, err := queryCoord.ReleasePartitions(ctx, &querypb.ReleasePartitionsRequest{
+			Base: &commonpb.MsgBase{
+				MsgType: commonpb.MsgType_ReleasePartitions,
+			},
+			CollectionID: defaultCollectionID,
+		})
+		assert.Equal(t, commonpb.ErrorCode_UnexpectedError, status.ErrorCode)
+		assert.NotNil(t, err)
+	})
+
+	t.Run("Test ReleaseNotExistPartition", func(t *testing.T) {
+		status, err := queryCoord.ReleasePartitions(ctx, &querypb.ReleasePartitionsRequest{
+			Base: &commonpb.MsgBase{
+				MsgType: commonpb.MsgType_ReleasePartitions,
+			},
+			CollectionID: defaultCollectionID,
+			PartitionIDs: []UniqueID{-1},
+		})
+		assert.Equal(t, commonpb.ErrorCode_Success, status.ErrorCode)
 		assert.Nil(t, err)
 	})
 
@@ -124,9 +218,19 @@ func TestGrpcTask(t *testing.T) {
 			CollectionID: defaultCollectionID,
 			PartitionIDs: []UniqueID{defaultPartitionID},
 		})
-		assert.Equal(t, status.ErrorCode, commonpb.ErrorCode_Success)
+		assert.Equal(t, commonpb.ErrorCode_Success, status.ErrorCode)
 		assert.Nil(t, err)
+	})
 
+	t.Run("Test ReleaseNotExistCollection", func(t *testing.T) {
+		status, err := queryCoord.ReleaseCollection(ctx, &querypb.ReleaseCollectionRequest{
+			Base: &commonpb.MsgBase{
+				MsgType: commonpb.MsgType_ReleaseCollection,
+			},
+			CollectionID: -1,
+		})
+		assert.Equal(t, commonpb.ErrorCode_Success, status.ErrorCode)
+		assert.Nil(t, err)
 	})
 
 	t.Run("Test ReleaseCollection", func(t *testing.T) {
@@ -136,7 +240,7 @@ func TestGrpcTask(t *testing.T) {
 			},
 			CollectionID: defaultCollectionID,
 		})
-		assert.Equal(t, status.ErrorCode, commonpb.ErrorCode_Success)
+		assert.Equal(t, commonpb.ErrorCode_Success, status.ErrorCode)
 		assert.Nil(t, err)
 	})
 
@@ -152,8 +256,8 @@ func TestGrpcTask(t *testing.T) {
 
 	t.Run("Test GetComponentStates", func(t *testing.T) {
 		states, err := queryCoord.GetComponentStates(ctx)
-		assert.Equal(t, states.Status.ErrorCode, commonpb.ErrorCode_Success)
-		assert.Equal(t, states.State.StateCode, internalpb.StateCode_Healthy)
+		assert.Equal(t, commonpb.ErrorCode_Success, states.Status.ErrorCode)
+		assert.Equal(t, internalpb.StateCode_Healthy, states.State.StateCode)
 		assert.Nil(t, err)
 	})
 
@@ -161,7 +265,7 @@ func TestGrpcTask(t *testing.T) {
 		res, err := queryCoord.CreateQueryChannel(ctx, &querypb.CreateQueryChannelRequest{
 			CollectionID: defaultCollectionID,
 		})
-		assert.Equal(t, res.Status.ErrorCode, commonpb.ErrorCode_Success)
+		assert.Equal(t, commonpb.ErrorCode_Success, res.Status.ErrorCode)
 		assert.Nil(t, err)
 	})
 
@@ -179,6 +283,32 @@ func TestGrpcTask(t *testing.T) {
 		assert.Equal(t, commonpb.ErrorCode_Success, res.Status.ErrorCode)
 	})
 
+	t.Run("Test InvalidMetricType", func(t *testing.T) {
+		metricReq := make(map[string]string)
+		metricReq["invalidKey"] = "invalidValue"
+		req, err := json.Marshal(metricReq)
+		assert.Nil(t, err)
+		res, err := queryCoord.GetMetrics(ctx, &milvuspb.GetMetricsRequest{
+			Base:    &commonpb.MsgBase{},
+			Request: string(req),
+		})
+
+		assert.NotNil(t, err)
+		assert.Equal(t, commonpb.ErrorCode_UnexpectedError, res.Status.ErrorCode)
+
+		metricReq = make(map[string]string)
+		metricReq[metricsinfo.MetricTypeKey] = "invalid"
+		req, err = json.Marshal(metricReq)
+		assert.Nil(t, err)
+		res, err = queryCoord.GetMetrics(ctx, &milvuspb.GetMetricsRequest{
+			Base:    &commonpb.MsgBase{},
+			Request: string(req),
+		})
+
+		assert.NotNil(t, err)
+		assert.Equal(t, commonpb.ErrorCode_UnexpectedError, res.Status.ErrorCode)
+	})
+
 	//nodes, err := queryCoord.cluster.getOnServiceNodes()
 	//assert.Nil(t, err)
 
@@ -257,3 +387,155 @@ func TestLoadBalanceTask(t *testing.T) {
 	queryNode2.stop()
 	queryCoord.Stop()
 }
+
+func TestGrpcTaskBeforeHealthy(t *testing.T) {
+	ctx := context.Background()
+	unHealthyCoord, err := startUnHealthyQueryCoord(ctx)
+	assert.Nil(t, err)
+
+	t.Run("Test LoadPartition", func(t *testing.T) {
+		status, err := unHealthyCoord.LoadPartitions(ctx, &querypb.LoadPartitionsRequest{
+			Base: &commonpb.MsgBase{
+				MsgType: commonpb.MsgType_LoadPartitions,
+			},
+			CollectionID: defaultCollectionID,
+			PartitionIDs: []UniqueID{defaultPartitionID},
+			Schema:       genCollectionSchema(defaultCollectionID, false),
+		})
+		assert.Equal(t, commonpb.ErrorCode_UnexpectedError, status.ErrorCode)
+		assert.NotNil(t, err)
+	})
+
+	t.Run("Test ShowPartitions", func(t *testing.T) {
+		res, err := unHealthyCoord.ShowPartitions(ctx, &querypb.ShowPartitionsRequest{
+			Base: &commonpb.MsgBase{
+				MsgType: commonpb.MsgType_ShowCollections,
+			},
+			CollectionID: defaultCollectionID,
+			PartitionIDs: []UniqueID{defaultPartitionID},
+		})
+		assert.Equal(t, commonpb.ErrorCode_UnexpectedError, res.Status.ErrorCode)
+		assert.NotNil(t, err)
+	})
+
+	t.Run("Test ShowAllPartitions", func(t *testing.T) {
+		res, err := unHealthyCoord.ShowPartitions(ctx, &querypb.ShowPartitionsRequest{
+			Base: &commonpb.MsgBase{
+				MsgType: commonpb.MsgType_ShowCollections,
+			},
+			CollectionID: defaultCollectionID,
+		})
+		assert.Equal(t, commonpb.ErrorCode_UnexpectedError, res.Status.ErrorCode)
+		assert.NotNil(t, err)
+	})
+
+	t.Run("Test GetPartitionStates", func(t *testing.T) {
+		res, err := unHealthyCoord.GetPartitionStates(ctx, &querypb.GetPartitionStatesRequest{
+			Base: &commonpb.MsgBase{
+				MsgType: commonpb.MsgType_GetPartitionStatistics,
+			},
+			CollectionID: defaultCollectionID,
+			PartitionIDs: []UniqueID{defaultPartitionID},
+		})
+		assert.Equal(t, commonpb.ErrorCode_UnexpectedError, res.Status.ErrorCode)
+		assert.NotNil(t, err)
+	})
+
+	t.Run("Test LoadCollection", func(t *testing.T) {
+		status, err := unHealthyCoord.LoadCollection(ctx, &querypb.LoadCollectionRequest{
+			Base: &commonpb.MsgBase{
+				MsgType: commonpb.MsgType_LoadCollection,
+			},
+			CollectionID: defaultCollectionID,
+			Schema:       genCollectionSchema(defaultCollectionID, false),
+		})
+		assert.Equal(t, commonpb.ErrorCode_UnexpectedError, status.ErrorCode)
+		assert.NotNil(t, err)
+	})
+
+	t.Run("Test ShowCollections", func(t *testing.T) {
+		res, err := unHealthyCoord.ShowCollections(ctx, &querypb.ShowCollectionsRequest{
+			Base: &commonpb.MsgBase{
+				MsgType: commonpb.MsgType_ShowCollections,
+			},
+			CollectionIDs: []UniqueID{defaultCollectionID},
+		})
+		assert.Equal(t, commonpb.ErrorCode_UnexpectedError, res.Status.ErrorCode)
+		assert.NotNil(t, err)
+	})
+
+	t.Run("Test ShowAllCollections", func(t *testing.T) {
+		res, err := unHealthyCoord.ShowCollections(ctx, &querypb.ShowCollectionsRequest{
+			Base: &commonpb.MsgBase{
+				MsgType: commonpb.MsgType_ShowCollections,
+			},
+		})
+		assert.Equal(t, commonpb.ErrorCode_UnexpectedError, res.Status.ErrorCode)
+		assert.NotNil(t, err)
+	})
+
+	t.Run("Test GetSegmentInfo", func(t *testing.T) {
+		res, err := unHealthyCoord.GetSegmentInfo(ctx, &querypb.GetSegmentInfoRequest{
+			Base: &commonpb.MsgBase{
+				MsgType: commonpb.MsgType_SegmentInfo,
+			},
+			SegmentIDs: []UniqueID{defaultSegmentID},
+		})
+		assert.Equal(t, commonpb.ErrorCode_UnexpectedError, res.Status.ErrorCode)
+		assert.NotNil(t, err)
+	})
+
+	t.Run("Test ReleasePartition", func(t *testing.T) {
+		status, err := unHealthyCoord.ReleasePartitions(ctx, &querypb.ReleasePartitionsRequest{
+			Base: &commonpb.MsgBase{
+				MsgType: commonpb.MsgType_ReleasePartitions,
+			},
+			CollectionID: defaultCollectionID,
+			PartitionIDs: []UniqueID{defaultPartitionID},
+		})
+		assert.Equal(t, commonpb.ErrorCode_UnexpectedError, status.ErrorCode)
+		assert.NotNil(t, err)
+
+	})
+
+	t.Run("Test ReleaseCollection", func(t *testing.T) {
+		status, err := unHealthyCoord.ReleaseCollection(ctx, &querypb.ReleaseCollectionRequest{
+			Base: &commonpb.MsgBase{
+				MsgType: commonpb.MsgType_ReleaseCollection,
+			},
+			CollectionID: defaultCollectionID,
+		})
+		assert.Equal(t, commonpb.ErrorCode_UnexpectedError, status.ErrorCode)
+		assert.NotNil(t, err)
+	})
+
+	t.Run("Test GetComponentStates", func(t *testing.T) {
+		states, err := unHealthyCoord.GetComponentStates(ctx)
+		assert.Equal(t, commonpb.ErrorCode_Success, states.Status.ErrorCode)
+		assert.Equal(t, internalpb.StateCode_Abnormal, states.State.StateCode)
+		assert.Nil(t, err)
+	})
+
+	t.Run("Test CreateQueryChannel", func(t *testing.T) {
+		res, err := unHealthyCoord.CreateQueryChannel(ctx, &querypb.CreateQueryChannelRequest{
+			CollectionID: defaultCollectionID,
+		})
+		assert.Equal(t, commonpb.ErrorCode_UnexpectedError, res.Status.ErrorCode)
+		assert.NotNil(t, err)
+	})
+
+	t.Run("Test GetMetrics", func(t *testing.T) {
+		metricReq := make(map[string]string)
+		metricReq[metricsinfo.MetricTypeKey] = "system_info"
+		req, err := json.Marshal(metricReq)
+		assert.Nil(t, err)
+		res, err := unHealthyCoord.GetMetrics(ctx, &milvuspb.GetMetricsRequest{
+			Base:    &commonpb.MsgBase{},
+			Request: string(req),
+		})
+
+		assert.Nil(t, err)
+		assert.Equal(t, commonpb.ErrorCode_UnexpectedError, res.Status.ErrorCode)
+	})
+
+}

internal/querycoord/querynode_test.go

@@ -59,6 +59,38 @@ func startQueryCoord(ctx context.Context) (*QueryCoord, error) {
 	return coord, nil
 }
 
+func startUnHealthyQueryCoord(ctx context.Context) (*QueryCoord, error) {
+	factory := msgstream.NewPmsFactory()
+
+	coord, err := NewQueryCoordTest(ctx, factory)
+	if err != nil {
+		return nil, err
+	}
+
+	rootCoord := newRootCoordMock()
+	rootCoord.createCollection(defaultCollectionID)
+	rootCoord.createPartition(defaultCollectionID, defaultPartitionID)
+
+	dataCoord, err := newDataCoordMock(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	coord.SetRootCoord(rootCoord)
+	coord.SetDataCoord(dataCoord)
+
+	err = coord.Register()
+	if err != nil {
+		return nil, err
+	}
+	err = coord.Init()
+	if err != nil {
+		return nil, err
+	}
+
+	return coord, nil
+}
+
 //func waitQueryNodeOnline(cluster *queryNodeCluster, nodeID int64)
 
 func waitAllQueryNodeOffline(cluster *queryNodeCluster, nodes map[int64]Node) bool {