From 06ad8e96be5fd516ce948b924b735686fa6c3878 Mon Sep 17 00:00:00 2001
From: Frank Schroeder <frank.schroeder@gmail.com>
Date: Mon, 26 Jun 2017 14:35:34 +0200
Subject: [PATCH] rpc: fix TestServer_Leave

wait for the leader election.
---
 agent/consul/server_test.go | 48 ++++++++++++++++++++-----------------
 1 file changed, 26 insertions(+), 22 deletions(-)

diff --git a/agent/consul/server_test.go b/agent/consul/server_test.go
index 84b79fce0..f57a6cab6 100644
--- a/agent/consul/server_test.go
+++ b/agent/consul/server_test.go
@@ -361,19 +361,21 @@ func TestServer_LeaveLeader(t *testing.T) {
 	// Try to join
 	joinLAN(t, s2, s1)
 
-	retry.Run(t, func(r *retry.R) {
-		r.Check(wantPeers(s1, 2))
-		r.Check(wantPeers(s2, 2))
-	})
+	testrpc.WaitForLeader(t, s1.RPC, "dc1")
+	testrpc.WaitForLeader(t, s2.RPC, "dc1")
 
 	// Issue a leave to the leader
-	for _, s := range []*Server{s1, s2} {
-		if !s.IsLeader() {
-			continue
-		}
-		if err := s.Leave(); err != nil {
-			t.Fatalf("err: %v", err)
-		}
+	var err error
+	switch {
+	case s1.IsLeader():
+		err = s1.Leave()
+	case s2.IsLeader():
+		err = s2.Leave()
+	default:
+		t.Fatal("no leader")
+	}
+	if err != nil {
+		t.Fatal("leave failed: ", err)
 	}
 
 	// Should lose a peer
@@ -396,19 +398,21 @@ func TestServer_Leave(t *testing.T) {
 	// Try to join
 	joinLAN(t, s2, s1)
 
-	retry.Run(t, func(r *retry.R) {
-		r.Check(wantPeers(s1, 2))
-		r.Check(wantPeers(s2, 2))
-	})
+	testrpc.WaitForLeader(t, s1.RPC, "dc1")
+	testrpc.WaitForLeader(t, s2.RPC, "dc1")
 
 	// Issue a leave to the non-leader
-	for _, s := range []*Server{s1, s2} {
-		if s.IsLeader() {
-			continue
-		}
-		if err := s.Leave(); err != nil {
-			t.Fatalf("err: %v", err)
-		}
+	var err error
+	switch {
+	case s1.IsLeader():
+		err = s2.Leave()
+	case s2.IsLeader():
+		err = s1.Leave()
+	default:
+		t.Fatal("no leader")
+	}
+	if err != nil {
+		t.Fatal("leave failed: ", err)
 	}
 
 	// Should lose a peer