From 16b35e07b31638ec1d227b47bccb5ed674a645d3 Mon Sep 17 00:00:00 2001 From: congqixia Date: Wed, 20 Sep 2023 18:59:23 +0800 Subject: [PATCH] Fix `TestSessionSuite/TestKeepAliveRetryActiveCancel` unit test logic (#27231) Signed-off-by: Congqi Xia --- internal/util/sessionutil/session_util.go | 2 +- internal/util/sessionutil/session_util_test.go | 8 +++----- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/internal/util/sessionutil/session_util.go b/internal/util/sessionutil/session_util.go index 5665f43e83..2788cc4eb2 100644 --- a/internal/util/sessionutil/session_util.go +++ b/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() } diff --git a/internal/util/sessionutil/session_util_test.go b/internal/util/sessionutil/session_util_test.go index 724dd9b2c1..e2e1801e21 100644 --- a/internal/util/sessionutil/session_util_test.go +++ b/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()) }