From f3e2d4a39a91b05eae171af21b37c11e10babf92 Mon Sep 17 00:00:00 2001 From: congqixia Date: Tue, 7 Mar 2023 09:51:50 +0800 Subject: [PATCH] Fix querynode stop meet unexpected GetComponentStates (#22590) Signed-off-by: Congqi Xia --- internal/querycoordv2/mocks/querynode.go | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/internal/querycoordv2/mocks/querynode.go b/internal/querycoordv2/mocks/querynode.go index 329bf7371b..30d6f3fc98 100644 --- a/internal/querycoordv2/mocks/querynode.go +++ b/internal/querycoordv2/mocks/querynode.go @@ -24,6 +24,7 @@ import ( "time" "github.com/milvus-io/milvus-proto/go-api/commonpb" + "github.com/milvus-io/milvus-proto/go-api/milvuspb" "github.com/milvus-io/milvus/internal/log" querypb "github.com/milvus-io/milvus/internal/proto/querypb" . "github.com/milvus-io/milvus/internal/querycoordv2/params" @@ -111,6 +112,22 @@ func (node *MockQueryNode) Start() error { segment.GetSegmentID()) node.segmentVersion[segment.GetSegmentID()] = req.GetVersion() }).Return(successStatus, nil).Maybe() + node.EXPECT().GetComponentStates(mock.Anything, mock.AnythingOfType("*milvuspb.GetComponentStatesRequest")). + Call.Return(func(context.Context, *milvuspb.GetComponentStatesRequest) *milvuspb.ComponentStates { + select { + case <-node.ctx.Done(): + return nil + default: + return &milvuspb.ComponentStates{} + } + }, func(context.Context, *milvuspb.GetComponentStatesRequest) error { + select { + case <-node.ctx.Done(): + return grpc.ErrServerStopped + default: + return nil + } + }).Maybe() // Register node.session.Init(typeutil.QueryNodeRole, node.addr, false, true)