From 6b5749723c81927ddeb2fb6323e4d6b8a99346af Mon Sep 17 00:00:00 2001
From: Simon Fell <fellforce@gmail.com>
Date: Mon, 18 Jan 2016 10:29:18 -0800
Subject: [PATCH] test LeaderRemoval while doing Apply's

---
 raft_test.go | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 deletions(-)

diff --git a/raft_test.go b/raft_test.go
index b4f7bfff843..2dd84a32018 100644
--- a/raft_test.go
+++ b/raft_test.go
@@ -251,7 +251,7 @@ CHECK:
 		if len(first.logs) != len(fsm.logs) {
 			fsm.Unlock()
 			if time.Now().After(limit) {
-				t.Fatalf("length mismatch: %d %d",
+				t.Fatalf("FSM log length mismatch: %d %d",
 					len(first.logs), len(fsm.logs))
 			} else {
 				goto WAIT
@@ -905,8 +905,18 @@ func TestRaft_RemoveLeader_NoShutdown(t *testing.T) {
 	leader := c.Leader()
 
 	// Remove the leader
-	leader.RemovePeer(leader.localAddr)
+	var removeFuture Future
+	for i := byte(0) ; i < 100; i++ {
+		leader.Apply([]byte{i}, 0)
+		if i == 80 {
+			removeFuture = leader.RemovePeer(leader.localAddr)
+		}
+	}
 
+	if err := removeFuture.Error(); err != nil {
+		t.Fatalf("RemovePeer failed with error %v", err)
+	}
+	
 	// Wait a while
 	time.Sleep(20 * time.Millisecond)
 
@@ -930,6 +940,9 @@ func TestRaft_RemoveLeader_NoShutdown(t *testing.T) {
 	if peers, _ := leader.peerStore.Peers(); len(peers) != 1 {
 		t.Fatalf("leader should have no peers")
 	}
+	
+	// Other nodes should have the same state
+	c.EnsureSame(t)
 }
 
 func TestRaft_RemoveLeader_SplitCluster(t *testing.T) {