From adfb5298c64954c1544090094e26d1007d18ef73 Mon Sep 17 00:00:00 2001
From: congqixia <congqi.xia@zilliz.com>
Date: Mon, 11 Sep 2023 18:13:17 +0800
Subject: [PATCH] Refine `TestSessionProcessActiveStandBy` unit test logic
 (#26980)

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
---
 internal/util/sessionutil/session_util_test.go | 13 ++++++-------
 1 file changed, 6 insertions(+), 7 deletions(-)

diff --git a/internal/util/sessionutil/session_util_test.go b/internal/util/sessionutil/session_util_test.go
index 9adda794e1..724dd9b2c1 100644
--- a/internal/util/sessionutil/session_util_test.go
+++ b/internal/util/sessionutil/session_util_test.go
@@ -603,8 +603,7 @@ func TestSessionProcessActiveStandBy(t *testing.T) {
 	defer etcdKV.RemoveWithPrefix("")
 
 	var wg sync.WaitGroup
-	ch := make(chan struct{})
-	signal := make(chan struct{}, 1)
+	signal := make(chan struct{})
 	flag := false
 
 	// register session 1, will be active
@@ -614,15 +613,15 @@ func TestSessionProcessActiveStandBy(t *testing.T) {
 	s1.SetEnableActiveStandBy(true)
 	s1.Register()
 	wg.Add(1)
-	s1.liveCh = ch
 	s1.ProcessActiveStandBy(func() error {
 		log.Debug("Session 1 become active")
 		wg.Done()
 		return nil
 	})
+	wg.Wait()
 	s1.LivenessCheck(ctx1, func() {
 		flag = true
-		signal <- struct{}{}
+		close(signal)
 		s1.cancelKeepAlive()
 	})
 	assert.False(t, s1.isStandby.Load().(bool))
@@ -645,14 +644,14 @@ func TestSessionProcessActiveStandBy(t *testing.T) {
 	// stop session 1, session 2 will take over primary service
 	log.Debug("Stop session 1, session 2 will take over primary service")
 	assert.False(t, flag)
-	ch <- struct{}{}
-	assert.False(t, flag)
-	s1.safeCloseLiveCh()
+
+	s1.Stop()
 	<-signal
 	assert.True(t, flag)
 
 	wg.Wait()
 	assert.False(t, s2.isStandby.Load().(bool))
+	s2.Stop()
 }
 
 func TestSessionEventType_String(t *testing.T) {