From 214711654e61deaa1f5ca77290b0a83db2e6c65d Mon Sep 17 00:00:00 2001
From: Lev Brouk <lev@synadia.com>
Date: Thu, 28 Sep 2023 12:42:18 -0700
Subject: [PATCH] PR feedback: use checkFor

---
 server/mqtt_test.go | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/server/mqtt_test.go b/server/mqtt_test.go
index 6ab76f19..5712be0e 100644
--- a/server/mqtt_test.go
+++ b/server/mqtt_test.go
@@ -4654,7 +4654,8 @@ func TestMQTTLockedSession(t *testing.T) {
 	s := testMQTTRunServer(t, o)
 	defer testMQTTShutdownServer(s)
 
-	ci := &mqttConnInfo{clientID: "sub", cleanSess: false}
+	subClientID := "sub"
+	ci := &mqttConnInfo{clientID: subClientID, cleanSess: false}
 	c, r := testMQTTConnect(t, ci, o.MQTT.Host, o.MQTT.Port)
 	defer c.Close()
 	testMQTTCheckConnAck(t, r, mqttConnAckRCConnectionAccepted, false)
@@ -4670,18 +4671,17 @@ func TestMQTTLockedSession(t *testing.T) {
 	// It is possible, however unlikely, to have received CONNACK while
 	// mqttProcessConnect is still running, and the session remains locked. Wait
 	// for it to finish.
-	for stillLocked := true; stillLocked; {
+	checkFor(t, 250*time.Millisecond, 10*time.Millisecond, func() error {
 		asm.mu.RLock()
-		_, stillLocked = asm.sessLocked["sub"]
-		asm.mu.RUnlock()
-
-		if stillLocked {
-			time.Sleep(1 * time.Millisecond)
+		defer asm.mu.RUnlock()
+		if _, stillLocked := asm.sessLocked[subClientID]; stillLocked {
+			return fmt.Errorf("session still locked")
 		}
-	}
+		return nil
+	})
 
 	// Get the session for "sub"
-	cli := testMQTTGetClient(t, s, "sub")
+	cli := testMQTTGetClient(t, s, subClientID)
 	sess := cli.mqtt.sess
 
 	// Pretend that the session above is locked.