De-flake tests (#5765)

This PR tries to deflake some unit tests.

Signed-off-by: Neil Twigg <neil@nats.io>

server/jetstream_cluster_1_test.go

@@ -6080,6 +6080,7 @@ func TestJetStreamClusterConsumerAckSyncReporting(t *testing.T) {
 
 	// Now ack the skipped message
 	skipped.AckSync()
+	c.waitOnAllCurrent()
 	for _, s := range c.servers {
 		jsz, err := s.Jsz(opts)
 		require_NoError(t, err)
@@ -6090,6 +6091,7 @@ func TestJetStreamClusterConsumerAckSyncReporting(t *testing.T) {
 
 	// Now ack the last message
 	last.AckSync()
+	c.waitOnAllCurrent()
 	for _, s := range c.servers {
 		jsz, err := s.Jsz(opts)
 		require_NoError(t, err)

server/jetstream_cluster_2_test.go

@@ -2660,22 +2660,15 @@ func TestJetStreamClusterStreamCatchupNoState(t *testing.T) {
 	nc.Close()
 	c.stopAll()
 	// Remove all state by truncating for the non-leader.
-	for _, fn := range []string{"1.blk", "1.idx", "1.fss"} {
-		fname := filepath.Join(config.StoreDir, "$G", "streams", "TEST", "msgs", fn)
-		fd, err := os.OpenFile(fname, os.O_RDWR, defaultFilePerms)
-		if err != nil {
-			continue
-		}
-		fd.Truncate(0)
-		fd.Close()
-	}
 	// For both make sure we have no raft snapshots.
 	snapDir := filepath.Join(lconfig.StoreDir, "$SYS", "_js_", gname, "snapshots")
-	os.RemoveAll(snapDir)
+	require_NoError(t, os.RemoveAll(snapDir))
+	msgsDir := filepath.Join(lconfig.StoreDir, "$SYS", "_js_", gname, "msgs")
+	require_NoError(t, os.RemoveAll(msgsDir))
 	// Remove all our raft state, we do not want to hold onto our term and index which
 	// results in a coin toss for who becomes the leader.
 	raftDir := filepath.Join(config.StoreDir, "$SYS", "_js_", gname)
-	os.RemoveAll(raftDir)
+	require_NoError(t, os.RemoveAll(raftDir))
 
 	// Now restart.
 	c.restartAll()

server/jetstream_helpers_test.go

@@ -1484,6 +1484,7 @@ func (c *cluster) waitOnAccount(account string) {
 func (c *cluster) waitOnClusterReady() {
 	c.t.Helper()
 	c.waitOnClusterReadyWithNumPeers(len(c.servers))
+	c.waitOnLeader()
 }
 
 func (c *cluster) waitOnClusterReadyWithNumPeers(numPeersExpected int) {

server/norace_test.go

@@ -10620,13 +10620,8 @@ func TestNoRaceJetStreamClusterMemoryStreamLastSequenceResetAfterRestart(t *test
 			s.Shutdown()
 			s.WaitForShutdown()
 			s = c.restartServer(s)
-			checkFor(t, 30*time.Second, time.Second, func() error {
-				hs := s.healthz(nil)
-				if hs.Error != _EMPTY_ {
-					return errors.New(hs.Error)
-				}
-				return nil
-			})
+			c.waitOnServerHealthz(s)
+			c.waitOnAllCurrent()
 			// Make sure all streams are current after healthz returns ok.
 			for i := 1; i <= numStreams; i++ {
 				stream := fmt.Sprintf("TEST:%d", i)
@@ -10693,13 +10688,8 @@ func TestNoRaceJetStreamClusterMemoryWorkQueueLastSequenceResetAfterRestart(t *t
 			s.Shutdown()
 			s.WaitForShutdown()
 			s = c.restartServer(s)
-			checkFor(t, 30*time.Second, time.Second, func() error {
-				hs := s.healthz(nil)
-				if hs.Error != _EMPTY_ {
-					return errors.New(hs.Error)
-				}
-				return nil
-			})
+			c.waitOnServerHealthz(s)
+			c.waitOnAllCurrent()
 			// Make sure all streams are current after healthz returns ok.
 			for i := 1; i <= numStreams; i++ {
 				stream := fmt.Sprintf("TEST:%d", i)

server/raft_test.go

@@ -360,23 +360,14 @@ func TestNRGSimpleElection(t *testing.T) {
 func TestNRGSwitchStateClearsQueues(t *testing.T) {
 	c := createJetStreamClusterExplicit(t, "R3S", 3)
 	defer c.shutdown()
+	s := c.servers[0] // RunBasicJetStreamServer not available
 
-	rg := c.createMemRaftGroup("TEST", 3, newStateAdder)
-	rg.waitOnLeader()
-
-	// Ensure there are no other nodes running that could
-	// send something into our IP queues or it may break the
-	// below assertions.
-	for _, n := range rg {
-		if !n.node().Leader() {
-			n.stop()
-		}
+	n := &raft{
+		prop:  newIPQueue[*Entry](s, "prop"),
+		resp:  newIPQueue[*appendEntryResponse](s, "resp"),
+		leadc: make(chan bool, 1), // for switchState
 	}
-
-	rg.lockAll()
-	defer rg.unlockAll()
-
-	n := rg.leader().node().(*raft)
+	n.state.Store(int32(Leader))
 	require_Equal(t, n.prop.len(), 0)
 	require_Equal(t, n.resp.len(), 0)