Fix `TestSessionSuite/TestKeepAliveRetryActiveCancel` unit test logic (#27231)

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>

internal/util/sessionutil/session_util.go

@@ -800,6 +800,7 @@ func (s *Session) LivenessCheck(ctx context.Context, callback func()) {
 	go func() {
 		defer s.wg.Done()
 		for {
+			defer s.SetDisconnected(true)
 			select {
 			case _, ok := <-s.liveCh:
 				// ok, still alive
@@ -808,7 +809,6 @@ func (s *Session) LivenessCheck(ctx context.Context, callback func()) {
 				}
 				// not ok, connection lost
 				log.Warn("connection lost detected, shuting down")
-				s.SetDisconnected(true)
 				if callback != nil {
 					go callback()
 				}

internal/util/sessionutil/session_util_test.go

@@ -919,17 +919,15 @@ func (s *SessionSuite) TestKeepAliveRetryActiveCancel() {
 
 	// Register
 	ch, err := session.registerService()
-	if err != nil {
-		panic(err)
-	}
+	s.Require().NoError(err)
 	session.liveCh = make(chan struct{})
 	session.processKeepAliveResponse(ch)
 	session.LivenessCheck(ctx, nil)
 	// active cancel, should not retry connect
 	session.cancelKeepAlive()
 
-	// sleep a while wait goroutine process
-	time.Sleep(time.Millisecond * 100)
+	// wait workers exit
+	session.wg.Wait()
 	// expected Disconnected = true, means session is closed
 	assert.Equal(s.T(), true, session.Disconnected())
 }