// Licensed to the LF AI & Data foundation under one // or more contributor license agreements. See the NOTICE file // distributed with this work for additional information // regarding copyright ownership. The ASF licenses this file // to you under the Apache License, Version 2.0 (the // "License"); you may not use this file except in compliance // with the License. You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package querycoord import ( "context" "encoding/json" "errors" "testing" "time" "github.com/milvus-io/milvus/internal/common" "github.com/milvus-io/milvus/internal/util/sessionutil" "github.com/stretchr/testify/assert" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/internalpb" "github.com/milvus-io/milvus/internal/proto/milvuspb" "github.com/milvus-io/milvus/internal/proto/querypb" "github.com/milvus-io/milvus/internal/util/metricsinfo" ) func waitLoadPartitionDone(ctx context.Context, queryCoord *QueryCoord, collectionID UniqueID, partitionIDs []UniqueID) error { for { showPartitionReq := &querypb.ShowPartitionsRequest{ Base: &commonpb.MsgBase{ MsgType: commonpb.MsgType_ShowPartitions, }, CollectionID: collectionID, PartitionIDs: partitionIDs, } res, err := queryCoord.ShowPartitions(ctx, showPartitionReq) if err != nil || res.Status.ErrorCode != commonpb.ErrorCode_Success { return errors.New("showPartitions failed") } loadDone := true for _, percent := range res.InMemoryPercentages { if percent < 100 { loadDone = false } } if loadDone { break } } return nil } func waitLoadCollectionDone(ctx context.Context, queryCoord *QueryCoord, collectionID UniqueID) error { for { showCollectionReq := &querypb.ShowCollectionsRequest{ Base: &commonpb.MsgBase{ MsgType: commonpb.MsgType_ShowPartitions, }, CollectionIDs: []UniqueID{collectionID}, } res, err := queryCoord.ShowCollections(ctx, showCollectionReq) if err != nil || res.Status.ErrorCode != commonpb.ErrorCode_Success { return errors.New("showCollection failed") } loadDone := true for _, percent := range res.InMemoryPercentages { if percent < 100 { loadDone = false } } if loadDone { break } } return nil } func TestGrpcTask(t *testing.T) { refreshParams() ctx := context.Background() queryCoord, err := startQueryCoord(ctx) assert.Nil(t, err) node, err := startQueryNodeServer(ctx) assert.Nil(t, err) waitQueryNodeOnline(queryCoord.cluster, node.queryNodeID) 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.Nil(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.Nil(t, err) }) t.Run("Test LoadPartition", 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) }) t.Run("Test ShowPartitions", 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_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.Nil(t, err) }) t.Run("Test ShowAllPartitions", func(t *testing.T) { res, err := queryCoord.ShowPartitions(ctx, &querypb.ShowPartitionsRequest{ Base: &commonpb.MsgBase{ MsgType: commonpb.MsgType_ShowCollections, }, CollectionID: defaultCollectionID, }) assert.Equal(t, commonpb.ErrorCode_Success, res.Status.ErrorCode) assert.Nil(t, err) }) t.Run("Test GetPartitionStates", func(t *testing.T) { res, err := queryCoord.GetPartitionStates(ctx, &querypb.GetPartitionStatesRequest{ Base: &commonpb.MsgBase{ MsgType: commonpb.MsgType_GetPartitionStatistics, }, CollectionID: defaultCollectionID, PartitionIDs: []UniqueID{defaultPartitionID}, }) assert.Equal(t, commonpb.ErrorCode_Success, res.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.Nil(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) }) t.Run("Test ReleasePartition", func(t *testing.T) { status, err := queryCoord.ReleasePartitions(ctx, &querypb.ReleasePartitionsRequest{ Base: &commonpb.MsgBase{ MsgType: commonpb.MsgType_ReleasePartitions, }, CollectionID: defaultCollectionID, PartitionIDs: []UniqueID{defaultPartitionID}, }) assert.Equal(t, commonpb.ErrorCode_Success, status.ErrorCode) assert.Nil(t, err) }) t.Run("Test LoadCollection", func(t *testing.T) { status, err := queryCoord.LoadCollection(ctx, &querypb.LoadCollectionRequest{ Base: &commonpb.MsgBase{ MsgType: commonpb.MsgType_LoadCollection, }, CollectionID: defaultCollectionID, Schema: genCollectionSchema(defaultCollectionID, false), }) 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) //}) t.Run("Test ShowCollections", func(t *testing.T) { res, err := queryCoord.ShowCollections(ctx, &querypb.ShowCollectionsRequest{ Base: &commonpb.MsgBase{ MsgType: commonpb.MsgType_ShowCollections, }, CollectionIDs: []UniqueID{defaultCollectionID}, }) 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.Nil(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, commonpb.ErrorCode_Success, res.Status.ErrorCode) assert.Nil(t, err) }) t.Run("Test GetSegmentInfo", func(t *testing.T) { res, err := queryCoord.GetSegmentInfo(ctx, &querypb.GetSegmentInfoRequest{ Base: &commonpb.MsgBase{ MsgType: commonpb.MsgType_SegmentInfo, }, SegmentIDs: []UniqueID{defaultSegmentID}, }) 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 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) { status, err := queryCoord.ReleaseCollection(ctx, &querypb.ReleaseCollectionRequest{ Base: &commonpb.MsgBase{ MsgType: commonpb.MsgType_ReleaseCollection, }, CollectionID: defaultCollectionID, }) assert.Equal(t, commonpb.ErrorCode_Success, status.ErrorCode) assert.Nil(t, err) }) t.Run("Test GetStatisticsChannel", func(t *testing.T) { _, err = queryCoord.GetStatisticsChannel(ctx) assert.Nil(t, err) }) t.Run("Test GetTimeTickChannel", func(t *testing.T) { _, err = queryCoord.GetTimeTickChannel(ctx) assert.Nil(t, err) }) t.Run("Test GetComponentStates", func(t *testing.T) { states, err := queryCoord.GetComponentStates(ctx) assert.Equal(t, commonpb.ErrorCode_Success, states.Status.ErrorCode) assert.Equal(t, internalpb.StateCode_Healthy, states.State.StateCode) assert.Nil(t, err) }) t.Run("Test CreateQueryChannel", func(t *testing.T) { res, err := queryCoord.CreateQueryChannel(ctx, &querypb.CreateQueryChannelRequest{ CollectionID: defaultCollectionID, }) assert.Equal(t, commonpb.ErrorCode_Success, res.Status.ErrorCode) assert.Nil(t, err) }) t.Run("Test LoadBalance", func(t *testing.T) { res, err := queryCoord.LoadBalance(ctx, &querypb.LoadBalanceRequest{ Base: &commonpb.MsgBase{ MsgType: commonpb.MsgType_LoadBalanceSegments, }, }) assert.Equal(t, commonpb.ErrorCode_UnexpectedError, res.ErrorCode) assert.Nil(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 := queryCoord.GetMetrics(ctx, &milvuspb.GetMetricsRequest{ Base: &commonpb.MsgBase{}, Request: string(req), }) assert.Nil(t, err) 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.Nil(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.Nil(t, err) assert.Equal(t, commonpb.ErrorCode_UnexpectedError, res.Status.ErrorCode) }) err = node.stop() err = removeNodeSession(node.queryNodeID) assert.Nil(t, err) queryCoord.Stop() err = removeAllSession() assert.Nil(t, err) } func TestGrpcTaskEnqueueFail(t *testing.T) { refreshParams() ctx := context.Background() queryCoord, err := startQueryCoord(ctx) assert.Nil(t, err) _, err = startQueryNodeServer(ctx) assert.Nil(t, err) taskIDAllocator := queryCoord.scheduler.taskIDAllocator failedAllocator := func() (UniqueID, error) { return 0, errors.New("scheduler failed to allocate ID") } queryCoord.scheduler.taskIDAllocator = failedAllocator t.Run("Test LoadPartition", 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_UnexpectedError, status.ErrorCode) assert.Nil(t, err) }) t.Run("Test LoadCollection", func(t *testing.T) { status, err := queryCoord.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.Nil(t, err) }) queryCoord.scheduler.taskIDAllocator = taskIDAllocator status, err := queryCoord.LoadCollection(ctx, &querypb.LoadCollectionRequest{ Base: &commonpb.MsgBase{ MsgType: commonpb.MsgType_LoadCollection, }, CollectionID: defaultCollectionID, Schema: genCollectionSchema(defaultCollectionID, false), }) assert.Equal(t, commonpb.ErrorCode_Success, status.ErrorCode) assert.Nil(t, err) queryCoord.scheduler.taskIDAllocator = failedAllocator t.Run("Test ReleaseCollection", func(t *testing.T) { status, err := queryCoord.ReleaseCollection(ctx, &querypb.ReleaseCollectionRequest{ Base: &commonpb.MsgBase{ MsgType: commonpb.MsgType_ReleaseCollection, }, CollectionID: defaultCollectionID, }) assert.Equal(t, commonpb.ErrorCode_UnexpectedError, status.ErrorCode) assert.Nil(t, err) }) queryCoord.scheduler.taskIDAllocator = taskIDAllocator status, err = queryCoord.ReleaseCollection(ctx, &querypb.ReleaseCollectionRequest{ Base: &commonpb.MsgBase{ MsgType: commonpb.MsgType_ReleaseCollection, }, CollectionID: defaultCollectionID, }) assert.Equal(t, commonpb.ErrorCode_Success, status.ErrorCode) assert.Nil(t, err) 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) queryCoord.scheduler.taskIDAllocator = failedAllocator t.Run("Test ReleasePartition", func(t *testing.T) { status, err := queryCoord.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.Nil(t, err) }) t.Run("Test LoadBalance", func(t *testing.T) { status, err := queryCoord.LoadBalance(ctx, &querypb.LoadBalanceRequest{ Base: &commonpb.MsgBase{ MsgType: commonpb.MsgType_ReleaseCollection, }, }) assert.NoError(t, err) assert.Equal(t, commonpb.ErrorCode_UnexpectedError, status.ErrorCode) }) queryCoord.Stop() err = removeAllSession() assert.Nil(t, err) } func TestLoadBalanceTask(t *testing.T) { refreshParams() baseCtx := context.Background() queryCoord, err := startQueryCoord(baseCtx) assert.Nil(t, err) queryNode1, err := startQueryNodeServer(baseCtx) assert.Nil(t, err) queryNode2, err := startQueryNodeServer(baseCtx) assert.Nil(t, err) time.Sleep(100 * time.Millisecond) res, err := queryCoord.LoadCollection(baseCtx, &querypb.LoadCollectionRequest{ Base: &commonpb.MsgBase{ MsgType: commonpb.MsgType_LoadCollection, }, CollectionID: defaultCollectionID, Schema: genCollectionSchema(defaultCollectionID, false), }) assert.Nil(t, err) assert.Equal(t, commonpb.ErrorCode_Success, res.ErrorCode) time.Sleep(100 * time.Millisecond) for { collectionInfo := queryCoord.meta.showCollections() if collectionInfo[0].InMemoryPercentage == 100 { break } } nodeID := queryNode1.queryNodeID queryCoord.cluster.stopNode(nodeID) loadBalanceSegment := &querypb.LoadBalanceRequest{ Base: &commonpb.MsgBase{ MsgType: commonpb.MsgType_LoadBalanceSegments, SourceID: nodeID, }, SourceNodeIDs: []int64{nodeID}, BalanceReason: querypb.TriggerCondition_NodeDown, } loadBalanceTask := &loadBalanceTask{ baseTask: &baseTask{ ctx: baseCtx, condition: newTaskCondition(baseCtx), triggerCondition: querypb.TriggerCondition_NodeDown, }, LoadBalanceRequest: loadBalanceSegment, rootCoord: queryCoord.rootCoordClient, dataCoord: queryCoord.dataCoordClient, indexCoord: queryCoord.indexCoordClient, cluster: queryCoord.cluster, meta: queryCoord.meta, } queryCoord.scheduler.Enqueue(loadBalanceTask) res, err = queryCoord.ReleaseCollection(baseCtx, &querypb.ReleaseCollectionRequest{ Base: &commonpb.MsgBase{ MsgType: commonpb.MsgType_ReleaseCollection, }, CollectionID: defaultCollectionID, }) assert.Nil(t, err) queryNode1.stop() queryNode2.stop() queryCoord.Stop() err = removeAllSession() assert.Nil(t, err) } func TestGrpcTaskBeforeHealthy(t *testing.T) { refreshParams() 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.Nil(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.Nil(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.Nil(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.Nil(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.Nil(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.Nil(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.Nil(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.Nil(t, err) }) t.Run("Test LoadBalance", func(t *testing.T) { res, err := unHealthyCoord.LoadBalance(ctx, &querypb.LoadBalanceRequest{ Base: &commonpb.MsgBase{ MsgType: commonpb.MsgType_LoadBalanceSegments, }, }) assert.NoError(t, err) assert.Equal(t, commonpb.ErrorCode_UnexpectedError, res.ErrorCode) }) 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.Nil(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.Nil(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.Nil(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) }) unHealthyCoord.Stop() err = removeAllSession() assert.Nil(t, err) } func Test_GrpcGetQueryChannelFail(t *testing.T) { kv := &testKv{ returnFn: failedResult, } meta, err := newMeta(context.Background(), kv, nil, nil) assert.Nil(t, err) queryCoord := &QueryCoord{ meta: meta, } queryCoord.stateCode.Store(internalpb.StateCode_Healthy) res, err := queryCoord.CreateQueryChannel(context.Background(), &querypb.CreateQueryChannelRequest{ CollectionID: defaultCollectionID, }) assert.Nil(t, err) assert.Equal(t, commonpb.ErrorCode_UnexpectedError, res.Status.ErrorCode) } func TestQueryCoord_GetComponentStates(t *testing.T) { n := &QueryCoord{} n.stateCode.Store(internalpb.StateCode_Healthy) resp, err := n.GetComponentStates(context.Background()) assert.NoError(t, err) assert.Equal(t, commonpb.ErrorCode_Success, resp.Status.ErrorCode) assert.Equal(t, common.NotRegisteredID, resp.State.NodeID) n.session = &sessionutil.Session{} n.session.UpdateRegistered(true) resp, err = n.GetComponentStates(context.Background()) assert.NoError(t, err) assert.Equal(t, commonpb.ErrorCode_Success, resp.Status.ErrorCode) } func Test_RepeatedLoadSameCollection(t *testing.T) { refreshParams() ctx := context.Background() queryCoord, err := startQueryCoord(ctx) assert.Nil(t, err) node, err := startQueryNodeServer(ctx) assert.Nil(t, err) waitQueryNodeOnline(queryCoord.cluster, node.queryNodeID) loadCollectionReq := &querypb.LoadCollectionRequest{ Base: &commonpb.MsgBase{ MsgType: commonpb.MsgType_LoadCollection, }, CollectionID: defaultCollectionID, Schema: genCollectionSchema(defaultCollectionID, false), } //first load defaultCollectionID status, err := queryCoord.LoadCollection(ctx, loadCollectionReq) assert.Equal(t, commonpb.ErrorCode_Success, status.ErrorCode) assert.Nil(t, err) waitLoadCollectionDone(ctx, queryCoord, defaultCollectionID) // second load defaultCollectionID status, err = queryCoord.LoadCollection(ctx, loadCollectionReq) assert.Equal(t, commonpb.ErrorCode_Success, status.ErrorCode) assert.Nil(t, err) node.stop() queryCoord.Stop() err = removeAllSession() assert.Nil(t, err) } func Test_LoadCollectionAndLoadPartitions(t *testing.T) { refreshParams() ctx := context.Background() queryCoord, err := startQueryCoord(ctx) assert.Nil(t, err) node, err := startQueryNodeServer(ctx) assert.Nil(t, err) waitQueryNodeOnline(queryCoord.cluster, node.queryNodeID) loadCollectionReq := &querypb.LoadCollectionRequest{ Base: &commonpb.MsgBase{ MsgType: commonpb.MsgType_LoadCollection, }, CollectionID: defaultCollectionID, Schema: genCollectionSchema(defaultCollectionID, false), } loadPartitionReq := &querypb.LoadPartitionsRequest{ Base: &commonpb.MsgBase{ MsgType: commonpb.MsgType_LoadPartitions, }, CollectionID: defaultCollectionID, PartitionIDs: []UniqueID{defaultPartitionID}, Schema: genCollectionSchema(defaultCollectionID, false), } //first load defaultCollectionID status, err := queryCoord.LoadCollection(ctx, loadCollectionReq) assert.Equal(t, commonpb.ErrorCode_Success, status.ErrorCode) assert.Nil(t, err) waitLoadCollectionDone(ctx, queryCoord, defaultCollectionID) // second load defaultPartitionID status, err = queryCoord.LoadPartitions(ctx, loadPartitionReq) assert.Equal(t, commonpb.ErrorCode_UnexpectedError, status.ErrorCode) assert.Nil(t, err) node.stop() queryCoord.Stop() err = removeAllSession() assert.Nil(t, err) } func Test_RepeatedLoadSamePartitions(t *testing.T) { refreshParams() ctx := context.Background() queryCoord, err := startQueryCoord(ctx) assert.Nil(t, err) node, err := startQueryNodeServer(ctx) assert.Nil(t, err) waitQueryNodeOnline(queryCoord.cluster, node.queryNodeID) loadPartitionReq := &querypb.LoadPartitionsRequest{ Base: &commonpb.MsgBase{ MsgType: commonpb.MsgType_LoadPartitions, }, CollectionID: defaultCollectionID, PartitionIDs: []UniqueID{defaultPartitionID}, Schema: genCollectionSchema(defaultCollectionID, false), } //first load defaultPartitionID status, err := queryCoord.LoadPartitions(ctx, loadPartitionReq) assert.Equal(t, commonpb.ErrorCode_Success, status.ErrorCode) assert.Nil(t, err) waitLoadPartitionDone(ctx, queryCoord, defaultCollectionID, []UniqueID{defaultPartitionID}) // second load defaultPartitionID status, err = queryCoord.LoadPartitions(ctx, loadPartitionReq) assert.Equal(t, commonpb.ErrorCode_Success, status.ErrorCode) assert.Nil(t, err) node.stop() queryCoord.Stop() err = removeAllSession() assert.Nil(t, err) } func Test_RepeatedLoadDifferentPartitions(t *testing.T) { refreshParams() ctx := context.Background() queryCoord, err := startQueryCoord(ctx) assert.Nil(t, err) node, err := startQueryNodeServer(ctx) assert.Nil(t, err) waitQueryNodeOnline(queryCoord.cluster, node.queryNodeID) loadPartitionReq := &querypb.LoadPartitionsRequest{ Base: &commonpb.MsgBase{ MsgType: commonpb.MsgType_LoadPartitions, }, CollectionID: defaultCollectionID, PartitionIDs: []UniqueID{defaultPartitionID}, Schema: genCollectionSchema(defaultCollectionID, false), } //first load defaultPartitionID status, err := queryCoord.LoadPartitions(ctx, loadPartitionReq) assert.Equal(t, commonpb.ErrorCode_Success, status.ErrorCode) assert.Nil(t, err) waitLoadPartitionDone(ctx, queryCoord, defaultCollectionID, []UniqueID{defaultPartitionID}) // second load defaultPartitionID+1 loadPartitionReq.PartitionIDs = []UniqueID{defaultPartitionID + 1} status, err = queryCoord.LoadPartitions(ctx, loadPartitionReq) assert.Equal(t, commonpb.ErrorCode_UnexpectedError, status.ErrorCode) assert.Nil(t, err) node.stop() queryCoord.Stop() err = removeAllSession() assert.Nil(t, err) } func Test_LoadPartitionsAndLoadCollection(t *testing.T) { refreshParams() ctx := context.Background() queryCoord, err := startQueryCoord(ctx) assert.Nil(t, err) node, err := startQueryNodeServer(ctx) assert.Nil(t, err) waitQueryNodeOnline(queryCoord.cluster, node.queryNodeID) loadCollectionReq := &querypb.LoadCollectionRequest{ Base: &commonpb.MsgBase{ MsgType: commonpb.MsgType_LoadCollection, }, CollectionID: defaultCollectionID, Schema: genCollectionSchema(defaultCollectionID, false), } loadPartitionReq := &querypb.LoadPartitionsRequest{ Base: &commonpb.MsgBase{ MsgType: commonpb.MsgType_LoadPartitions, }, CollectionID: defaultCollectionID, PartitionIDs: []UniqueID{defaultPartitionID}, Schema: genCollectionSchema(defaultCollectionID, false), } //first load defaultPartitionID status, err := queryCoord.LoadPartitions(ctx, loadPartitionReq) assert.Equal(t, commonpb.ErrorCode_Success, status.ErrorCode) assert.Nil(t, err) waitLoadPartitionDone(ctx, queryCoord, defaultCollectionID, []UniqueID{defaultPartitionID}) // second load defaultCollectionID status, err = queryCoord.LoadCollection(ctx, loadCollectionReq) assert.Equal(t, commonpb.ErrorCode_UnexpectedError, status.ErrorCode) assert.Nil(t, err) node.stop() queryCoord.Stop() err = removeAllSession() assert.Nil(t, err) } func Test_LoadAndReleaseCollection(t *testing.T) { refreshParams() ctx := context.Background() queryCoord, err := startQueryCoord(ctx) assert.Nil(t, err) node, err := startQueryNodeServer(ctx) assert.Nil(t, err) waitQueryNodeOnline(queryCoord.cluster, node.queryNodeID) loadCollectionReq := &querypb.LoadCollectionRequest{ Base: &commonpb.MsgBase{ MsgType: commonpb.MsgType_LoadCollection, }, CollectionID: defaultCollectionID, Schema: genCollectionSchema(defaultCollectionID, false), } releaseCollectionReq := &querypb.ReleaseCollectionRequest{ Base: &commonpb.MsgBase{ MsgType: commonpb.MsgType_ReleaseCollection, }, CollectionID: defaultCollectionID, } //first load defaultCollectionID status, err := queryCoord.LoadCollection(ctx, loadCollectionReq) assert.Equal(t, commonpb.ErrorCode_Success, status.ErrorCode) assert.Nil(t, err) waitLoadCollectionDone(ctx, queryCoord, defaultCollectionID) // second release defaultCollectionID status, err = queryCoord.ReleaseCollection(ctx, releaseCollectionReq) assert.Equal(t, commonpb.ErrorCode_Success, status.ErrorCode) assert.Nil(t, err) node.stop() queryCoord.Stop() err = removeAllSession() assert.Nil(t, err) } func Test_LoadAndReleasePartitions(t *testing.T) { refreshParams() ctx := context.Background() queryCoord, err := startQueryCoord(ctx) assert.Nil(t, err) node, err := startQueryNodeServer(ctx) assert.Nil(t, err) waitQueryNodeOnline(queryCoord.cluster, node.queryNodeID) loadPartitionReq := &querypb.LoadPartitionsRequest{ Base: &commonpb.MsgBase{ MsgType: commonpb.MsgType_LoadPartitions, }, CollectionID: defaultCollectionID, PartitionIDs: []UniqueID{defaultPartitionID}, Schema: genCollectionSchema(defaultCollectionID, false), } releasePartitionReq := &querypb.ReleasePartitionsRequest{ Base: &commonpb.MsgBase{ MsgType: commonpb.MsgType_ReleasePartitions, }, CollectionID: defaultCollectionID, PartitionIDs: []UniqueID{defaultPartitionID}, } //first load defaultPartitionID status, err := queryCoord.LoadPartitions(ctx, loadPartitionReq) assert.Equal(t, commonpb.ErrorCode_Success, status.ErrorCode) assert.Nil(t, err) waitLoadPartitionDone(ctx, queryCoord, defaultCollectionID, []UniqueID{defaultPartitionID}) // second release defaultPartitionID status, err = queryCoord.ReleasePartitions(ctx, releasePartitionReq) assert.Equal(t, commonpb.ErrorCode_Success, status.ErrorCode) assert.Nil(t, err) node.stop() queryCoord.Stop() err = removeAllSession() assert.Nil(t, err) } func Test_LoadCollectionAndReleasePartitions(t *testing.T) { refreshParams() ctx := context.Background() queryCoord, err := startQueryCoord(ctx) assert.Nil(t, err) node, err := startQueryNodeServer(ctx) assert.Nil(t, err) waitQueryNodeOnline(queryCoord.cluster, node.queryNodeID) loadCollectionReq := &querypb.LoadCollectionRequest{ Base: &commonpb.MsgBase{ MsgType: commonpb.MsgType_LoadCollection, }, CollectionID: defaultCollectionID, Schema: genCollectionSchema(defaultCollectionID, false), } releasePartitionReq := &querypb.ReleasePartitionsRequest{ Base: &commonpb.MsgBase{ MsgType: commonpb.MsgType_ReleasePartitions, }, CollectionID: defaultCollectionID, PartitionIDs: []UniqueID{defaultPartitionID}, } //first load defaultCollectionID status, err := queryCoord.LoadCollection(ctx, loadCollectionReq) assert.Equal(t, commonpb.ErrorCode_Success, status.ErrorCode) assert.Nil(t, err) waitLoadCollectionDone(ctx, queryCoord, defaultCollectionID) // second release defaultPartitionID status, err = queryCoord.ReleasePartitions(ctx, releasePartitionReq) assert.Equal(t, commonpb.ErrorCode_UnexpectedError, status.ErrorCode) assert.Nil(t, err) node.stop() queryCoord.Stop() err = removeAllSession() assert.Nil(t, err) } func Test_LoadPartitionsAndReleaseCollection(t *testing.T) { refreshParams() ctx := context.Background() queryCoord, err := startQueryCoord(ctx) assert.Nil(t, err) node, err := startQueryNodeServer(ctx) assert.Nil(t, err) waitQueryNodeOnline(queryCoord.cluster, node.queryNodeID) loadPartitionReq := &querypb.LoadPartitionsRequest{ Base: &commonpb.MsgBase{ MsgType: commonpb.MsgType_LoadPartitions, }, CollectionID: defaultCollectionID, PartitionIDs: []UniqueID{defaultPartitionID}, Schema: genCollectionSchema(defaultCollectionID, false), } releaseCollectionReq := &querypb.ReleaseCollectionRequest{ Base: &commonpb.MsgBase{ MsgType: commonpb.MsgType_ReleaseCollection, }, CollectionID: defaultCollectionID, } //first load defaultPartitionID status, err := queryCoord.LoadPartitions(ctx, loadPartitionReq) assert.Equal(t, commonpb.ErrorCode_Success, status.ErrorCode) assert.Nil(t, err) waitLoadPartitionDone(ctx, queryCoord, defaultCollectionID, []UniqueID{defaultPartitionID}) // second release defaultCollectionID status, err = queryCoord.ReleaseCollection(ctx, releaseCollectionReq) assert.Equal(t, commonpb.ErrorCode_Success, status.ErrorCode) assert.Nil(t, err) node.stop() queryCoord.Stop() err = removeAllSession() assert.Nil(t, err) } func Test_RepeatedReleaseCollection(t *testing.T) { refreshParams() ctx := context.Background() queryCoord, err := startQueryCoord(ctx) assert.Nil(t, err) node, err := startQueryNodeServer(ctx) assert.Nil(t, err) waitQueryNodeOnline(queryCoord.cluster, node.queryNodeID) loadCollectionReq := &querypb.LoadCollectionRequest{ Base: &commonpb.MsgBase{ MsgType: commonpb.MsgType_LoadCollection, }, CollectionID: defaultCollectionID, Schema: genCollectionSchema(defaultCollectionID, false), } releaseCollectionReq := &querypb.ReleaseCollectionRequest{ Base: &commonpb.MsgBase{ MsgType: commonpb.MsgType_ReleaseCollection, }, CollectionID: defaultCollectionID, } // load defaultCollectionID status, err := queryCoord.LoadCollection(ctx, loadCollectionReq) assert.Equal(t, commonpb.ErrorCode_Success, status.ErrorCode) assert.Nil(t, err) waitLoadCollectionDone(ctx, queryCoord, defaultCollectionID) // first release defaultCollectionID status, err = queryCoord.ReleaseCollection(ctx, releaseCollectionReq) assert.Equal(t, commonpb.ErrorCode_Success, status.ErrorCode) assert.Nil(t, err) // second release defaultCollectionID status, err = queryCoord.ReleaseCollection(ctx, releaseCollectionReq) assert.Equal(t, commonpb.ErrorCode_Success, status.ErrorCode) assert.Nil(t, err) node.stop() queryCoord.Stop() err = removeAllSession() assert.Nil(t, err) } func Test_RepeatedReleaseSamePartitions(t *testing.T) { refreshParams() ctx := context.Background() queryCoord, err := startQueryCoord(ctx) assert.Nil(t, err) node, err := startQueryNodeServer(ctx) assert.Nil(t, err) waitQueryNodeOnline(queryCoord.cluster, node.queryNodeID) loadPartitionReq := &querypb.LoadPartitionsRequest{ Base: &commonpb.MsgBase{ MsgType: commonpb.MsgType_LoadPartitions, }, CollectionID: defaultCollectionID, PartitionIDs: []UniqueID{defaultPartitionID}, Schema: genCollectionSchema(defaultCollectionID, false), } releasePartitionReq := &querypb.ReleasePartitionsRequest{ Base: &commonpb.MsgBase{ MsgType: commonpb.MsgType_ReleasePartitions, }, CollectionID: defaultCollectionID, PartitionIDs: []UniqueID{defaultPartitionID}, } // load defaultPartitionID status, err := queryCoord.LoadPartitions(ctx, loadPartitionReq) assert.Equal(t, commonpb.ErrorCode_Success, status.ErrorCode) assert.Nil(t, err) waitLoadPartitionDone(ctx, queryCoord, defaultCollectionID, []UniqueID{defaultPartitionID}) // first release defaultPartitionID status, err = queryCoord.ReleasePartitions(ctx, releasePartitionReq) assert.Equal(t, commonpb.ErrorCode_Success, status.ErrorCode) assert.Nil(t, err) // second release defaultPartitionID status, err = queryCoord.ReleasePartitions(ctx, releasePartitionReq) assert.Equal(t, commonpb.ErrorCode_Success, status.ErrorCode) assert.Nil(t, err) node.stop() queryCoord.Stop() err = removeAllSession() assert.Nil(t, err) } func Test_RepeatedReleaseDifferentPartitions(t *testing.T) { refreshParams() ctx := context.Background() queryCoord, err := startQueryCoord(ctx) assert.Nil(t, err) node, err := startQueryNodeServer(ctx) assert.Nil(t, err) waitQueryNodeOnline(queryCoord.cluster, node.queryNodeID) loadPartitionReq := &querypb.LoadPartitionsRequest{ Base: &commonpb.MsgBase{ MsgType: commonpb.MsgType_LoadPartitions, }, CollectionID: defaultCollectionID, PartitionIDs: []UniqueID{defaultPartitionID, defaultPartitionID + 1}, Schema: genCollectionSchema(defaultCollectionID, false), } releasePartitionReq := &querypb.ReleasePartitionsRequest{ Base: &commonpb.MsgBase{ MsgType: commonpb.MsgType_ReleasePartitions, }, CollectionID: defaultCollectionID, PartitionIDs: []UniqueID{defaultPartitionID}, } // load defaultPartitionID and defaultPartitionID+1 status, err := queryCoord.LoadPartitions(ctx, loadPartitionReq) assert.Equal(t, commonpb.ErrorCode_Success, status.ErrorCode) assert.Nil(t, err) waitLoadPartitionDone(ctx, queryCoord, defaultCollectionID, []UniqueID{defaultPartitionID, defaultPartitionID + 1}) // first release defaultPartitionID status, err = queryCoord.ReleasePartitions(ctx, releasePartitionReq) assert.Equal(t, commonpb.ErrorCode_Success, status.ErrorCode) assert.Nil(t, err) // second release defaultPartitionID+1 releasePartitionReq.PartitionIDs = []UniqueID{defaultPartitionID + 1} status, err = queryCoord.ReleasePartitions(ctx, releasePartitionReq) assert.Equal(t, commonpb.ErrorCode_Success, status.ErrorCode) assert.Nil(t, err) node.stop() queryCoord.Stop() err = removeAllSession() assert.Nil(t, err) }