From 9b7c81c37eeeb42cbee43162c78f02a1eef80959 Mon Sep 17 00:00:00 2001
From: Derek Collison <derek@nats.io>
Date: Sun, 3 Jul 2022 12:45:27 -0700
Subject: [PATCH] Some tests improvements on non-standard JS cluster setup

Signed-off-by: Derek Collison <derek@nats.io>
---
 server/jetstream_helpers_test.go |  2 ++
 server/monitor_test.go           | 15 ++++-----------
 server/mqtt_test.go              |  6 ++++--
 3 files changed, 10 insertions(+), 13 deletions(-)

diff --git a/server/jetstream_helpers_test.go b/server/jetstream_helpers_test.go
index b8c2ff25..3124bfa3 100644
--- a/server/jetstream_helpers_test.go
+++ b/server/jetstream_helpers_test.go
@@ -1257,11 +1257,13 @@ func (c *cluster) waitOnClusterReadyWithNumPeers(numPeersExpected int) {
 	}
 
 	if leader == nil {
+		c.shutdown()
 		c.t.Fatalf("Failed to elect a meta-leader")
 	}
 
 	peersSeen := len(leader.JetStreamClusterPeers())
 	c.shutdown()
+
 	if leader == nil {
 		c.t.Fatalf("Expected a cluster leader and fully formed cluster, no leader")
 	} else {
diff --git a/server/monitor_test.go b/server/monitor_test.go
index e4b6bc30..9a0db9ed 100644
--- a/server/monitor_test.go
+++ b/server/monitor_test.go
@@ -3984,17 +3984,10 @@ func TestMonitorAuthorizedUsers(t *testing.T) {
 // Helper function to check that a JS cluster is formed
 func checkForJSClusterUp(t *testing.T, servers ...*Server) {
 	t.Helper()
-	checkFor(t, 10*time.Second, 100*time.Millisecond, func() error {
-		for _, s := range servers {
-			if !s.JetStreamEnabled() {
-				return fmt.Errorf("jetstream not enabled")
-			}
-			if !s.JetStreamIsCurrent() {
-				return fmt.Errorf("jetstream not current")
-			}
-		}
-		return nil
-	})
+	// We will use the other JetStream helpers here.
+	c := &cluster{t: t, servers: servers}
+	c.checkClusterFormed()
+	c.waitOnClusterReady()
 }
 
 func TestMonitorJszNonJszServer(t *testing.T) {
diff --git a/server/mqtt_test.go b/server/mqtt_test.go
index 91decee0..f7fce85a 100644
--- a/server/mqtt_test.go
+++ b/server/mqtt_test.go
@@ -3167,8 +3167,10 @@ func TestMQTTLeafnodeWithoutJSToClusterWithJSNoSharedSysAcc(t *testing.T) {
 			for _, s := range cluster {
 				if s.JetStreamIsLeader() {
 					// Need to wait for usage updates now to propagate to meta leader.
-					time.Sleep(250 * time.Millisecond)
-					return nil
+					if len(s.JetStreamClusterPeers()) == len(cluster) {
+						time.Sleep(100 * time.Millisecond)
+						return nil
+					}
 				}
 			}
 			return fmt.Errorf("no leader yet")