fix ut

Signed-off-by: Cai Zhang <[email protected]>
milvus-io · Dec 18, 2024 · c9d082f · c9d082f
1 parent 457113f
commit c9d082f
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 8 deletions.
diff --git 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
@@ -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:     "",