De-flake RAFT state adder tests

Signed-off-by: Maurice van Veen <github@mauricevanveen.com>

server/raft_helpers_test.go

@@ -246,6 +246,7 @@ func (a *stateAdder) stop() {
 	a.Lock()
 	defer a.Unlock()
 	a.n.Stop()
+	a.n.WaitForStop()
 }
 
 // Restart the group
@@ -273,6 +274,11 @@ func (a *stateAdder) restart() {
 		panic(err)
 	}
 
+	// Must reset in-memory state.
+	// A real restart would not preserve it, but more importantly we have no way to detect if we
+	// already applied an entry. So, the sum must only be updated based on append entries or snapshots.
+	a.sum = 0
+
 	a.n, err = a.s.startRaftNode(globalAccountName, a.cfg, pprofLabels{})
 	if err != nil {
 		panic(err)

server/raft_test.go

@@ -605,8 +605,9 @@ func TestNRGHeartbeatOnLeaderChange(t *testing.T) {
 		leader := rg.leader().(*stateAdder)
 		leader.proposeDelta(22)
 		leader.proposeDelta(-11)
-		leader.proposeDelta(-11)
-		rg.waitOnTotal(t, 0)
+		leader.proposeDelta(-10)
+		// Must observe forward progress, so each iteration will check +1 total.
+		rg.waitOnTotal(t, int64(i+1))
 		leader.stop()
 		leader.restart()
 		rg.waitOnLeader()