add wait loop to make test more robust

raft_test.go

@@ -2379,6 +2379,7 @@ func TestRaft_RemovedFollower_Vote(t *testing.T) {
 	if configuration := c.getConfiguration(followers[1]); len(configuration.Servers) != 2 {
 		t.Fatalf("too many peers")
 	}
+	waitforState(followerRemoved, Follower)
 	// The removed node should be still in Follower state
 	require.Equal(t, Follower, followerRemoved.getState())
 
@@ -2401,11 +2402,7 @@ func TestRaft_RemovedFollower_Vote(t *testing.T) {
 	time.Sleep(c.propagateTimeout)
 
 	// wait for the remaining follower to trigger an election
-	count := 0
-	for follower.getState() != Candidate && count < 1000 {
-		count++
-		time.Sleep(1 * time.Millisecond)
-	}
+	waitforState(follower, Candidate)
 	require.Equal(t, Candidate, follower.getState())
 
 	// send a vote request from the removed follower to the Candidate follower
@@ -2418,3 +2415,11 @@ func TestRaft_RemovedFollower_Vote(t *testing.T) {
 		t.Fatalf("expected vote to not be granted, but it was %+v", resp)
 	}
 }
+
+func waitforState(follower *Raft, state RaftState) {
+	count := 0
+	for follower.getState() != state && count < 1000 {
+		count++
+		time.Sleep(1 * time.Millisecond)
+	}
+}