fix flake in test for restarted followers not causing an election

raft_test.go

@@ -3106,14 +3106,37 @@ func TestRaft_FollowerRemovalNoElection(t *testing.T) {
 	if err != nil {
 		t.Fatalf("error restarting follower: %v", err)
 	}
+	n.RegisterObserver(NewObserver(c.observationCh, false, nil))
+
 	c.rafts[i] = n
 	c.trans[i] = n.trans.(*InmemTransport)
 	c.fsms[i] = n.fsm.(*MockFSM)
+	time.Sleep(500 * time.Millisecond)
 	c.FullyConnect()
-	// There should be no re-election during this sleep
-	time.Sleep(250 * time.Millisecond)
 
-	// Let things settle and make sure we recovered.
+	// Wait for the restarted follower to settle.
+	ch := c.WaitEventChan(t.Context(), func(o *Observation) bool {
+		if o == nil {
+			return false
+		}
+		if o.Raft.localID == follower.localID {
+			if obs, ok := o.Data.(LeaderObservation); ok {
+				if obs.Leader != "" {
+					t.Logf("[INFO] restarted follower has rejoined cluster")
+					return true
+				}
+			}
+
+		}
+		return false
+	})
+	select {
+	case <-time.After(time.Second):
+		t.Fatal("restarted follower never got leader")
+	case <-ch:
+	}
+
+	// There should have been no re-election
 	c.EnsureLeader(t, leader.localAddr)
 	c.EnsureSame(t)
 	c.EnsureSamePeers(t)

testing.go

@@ -587,6 +587,7 @@ func (c *cluster) IndexOf(r *Raft) int {
 // EnsureLeader checks that ALL the nodes think the leader is the given expected
 // leader.
 func (c *cluster) EnsureLeader(t *testing.T, expect ServerAddress) {
+	t.Helper()
 	// We assume c.Leader() has been called already; now check all the rafts
 	// think the leader is correct
 	fail := false
@@ -612,6 +613,7 @@ func (c *cluster) EnsureLeader(t *testing.T, expect ServerAddress) {
 
 // EnsureSame makes sure all the FSMs have the same contents.
 func (c *cluster) EnsureSame(t *testing.T) {
+	t.Helper()
 	limit := time.Now().Add(c.longstopTimeout)
 	first := getMockFSM(c.fsms[0])
 
@@ -690,6 +692,7 @@ func (c *cluster) getConfiguration(r *Raft) Configuration {
 
 // EnsureSamePeers makes sure all the rafts have the same set of peers.
 func (c *cluster) EnsureSamePeers(t *testing.T) {
+	t.Helper()
 	limit := time.Now().Add(c.longstopTimeout)
 	peerSet := c.getConfiguration(c.rafts[0])