From 50cecc65ed725810d72382cee8a9981de3a57eee Mon Sep 17 00:00:00 2001
From: congqixia <congqi.xia@zilliz.com>
Date: Sat, 18 Jun 2022 18:22:12 +0800
Subject: [PATCH] Fix mock querynode server session not revoked (#17229)

Revoke mock querynode server session when it's stopped
This PR reduces the running time of TestLoadBalanceIndexedSegmentsAfterNodeDown from 60+ seconds to less than 1+ seconds
Also related to #17212 #17215

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
---
 internal/querycoord/mock_querynode_server_test.go | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/internal/querycoord/mock_querynode_server_test.go b/internal/querycoord/mock_querynode_server_test.go
index c54fbd2c80..4dcdae5e13 100644
--- a/internal/querycoord/mock_querynode_server_test.go
+++ b/internal/querycoord/mock_querynode_server_test.go
@@ -22,6 +22,7 @@ import (
 	"net"
 	"strconv"
 	"sync"
+	"time"
 
 	"go.uber.org/zap"
 	"google.golang.org/grpc"
@@ -172,6 +173,9 @@ func (qs *queryNodeServerMock) start() error {
 
 func (qs *queryNodeServerMock) stop() error {
 	qs.cancel()
+	if qs.session != nil {
+		qs.session.Revoke(time.Second)
+	}
 	if qs.grpcServer != nil {
 		qs.grpcServer.GracefulStop()
 	}