From c9d082f102d953f85c1b69aa189ca085c39dca52 Mon Sep 17 00:00:00 2001
From: Cai Zhang <cai.zhang@zilliz.com>
Date: Wed, 18 Dec 2024 13:06:53 +0800
Subject: [PATCH] fix ut

Signed-off-by: Cai Zhang <cai.zhang@zilliz.com>
---
 internal/datacoord/handler.go      |  3 +++
 internal/datacoord/handler_test.go | 20 ++++++++++++--------
 2 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/internal/datacoord/handler.go b/internal/datacoord/handler.go
index 9b44352590187..93a4657180245 100644
--- a/internal/datacoord/handler.go
+++ b/internal/datacoord/handler.go
@@ -26,6 +26,7 @@ import (
 	"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
 	"github.com/milvus-io/milvus-proto/go-api/v2/msgpb"
 	"github.com/milvus-io/milvus/internal/proto/datapb"
+	"github.com/milvus-io/milvus/pkg/common"
 	"github.com/milvus-io/milvus/pkg/log"
 	"github.com/milvus-io/milvus/pkg/util/funcutil"
 	"github.com/milvus-io/milvus/pkg/util/retry"
@@ -127,6 +128,7 @@ func (h *ServerHandler) GetQueryVChanPositions(channel RWChannel, partitionIDs .
 		validPartitions = collInfo.Partitions
 	}
 	validPartitionsMap := make(map[int64]bool)
+	validPartitionsMap[common.AllPartitionsID] = true
 	for _, partitionID := range validPartitions {
 		currentPartitionStatsVersion := h.s.meta.partitionStatsMeta.GetCurrentPartitionStatsVersion(channel.GetCollectionID(), partitionID, channel.GetName())
 		partStatsVersionsMap[partitionID] = currentPartitionStatsVersion
@@ -296,6 +298,7 @@ func (h *ServerHandler) GetCurrentSegmentsView(ctx context.Context, channel RWCh
 		validPartitions = collInfo.Partitions
 	}
 	validPartitionsMap := make(map[int64]bool)
+	validPartitionsMap[common.AllPartitionsID] = true
 	for _, partitionID := range validPartitions {
 		validPartitionsMap[partitionID] = true
 	}
diff --git a/internal/datacoord/handler_test.go b/internal/datacoord/handler_test.go
index 1535f3c02f3fe..5ef59b62e8a39 100644
--- a/internal/datacoord/handler_test.go
+++ b/internal/datacoord/handler_test.go
@@ -30,8 +30,9 @@ func TestGetQueryVChanPositionsRetrieveM2N(t *testing.T) {
 
 	channel := "ch1"
 	svr.meta.AddCollection(&collectionInfo{
-		ID:     1,
-		Schema: schema,
+		ID:         1,
+		Partitions: []int64{0},
+		Schema:     schema,
 		StartPositions: []*commonpb.KeyDataPair{
 			{
 				Key:  channel,
@@ -130,8 +131,9 @@ func TestGetQueryVChanPositions(t *testing.T) {
 	defer closeTestServer(t, svr)
 	schema := newTestSchema()
 	svr.meta.AddCollection(&collectionInfo{
-		ID:     0,
-		Schema: schema,
+		ID:         0,
+		Partitions: []int64{0, 1},
+		Schema:     schema,
 		StartPositions: []*commonpb.KeyDataPair{
 			{
 				Key:  "ch1",
@@ -583,8 +585,9 @@ func TestGetQueryVChanPositions_Retrieve_unIndexed(t *testing.T) {
 		defer closeTestServer(t, svr)
 		schema := newTestSchema()
 		svr.meta.AddCollection(&collectionInfo{
-			ID:     0,
-			Schema: schema,
+			ID:         0,
+			Partitions: []int64{0},
+			Schema:     schema,
 		})
 		err := svr.meta.indexMeta.CreateIndex(context.TODO(), &model.Index{
 			TenantID:     "",
@@ -964,8 +967,9 @@ func TestGetCurrentSegmentsView(t *testing.T) {
 	defer closeTestServer(t, svr)
 	schema := newTestSchema()
 	svr.meta.AddCollection(&collectionInfo{
-		ID:     0,
-		Schema: schema,
+		ID:         0,
+		Partitions: []int64{0},
+		Schema:     schema,
 	})
 	err := svr.meta.indexMeta.CreateIndex(context.TODO(), &model.Index{
 		TenantID:     "",