From e0694185140207ed63b2887769ef15e35ffd74e5 Mon Sep 17 00:00:00 2001
From: Pierre Souchay <p.souchay@criteo.com>
Date: Thu, 23 Aug 2018 09:12:59 +0200
Subject: [PATCH] More stable Unit Test / fixes according to @freddygv
 suggestions

---
 agent/consul/catalog_endpoint_test.go | 13 +++++--------
 testrpc/wait.go                       | 14 ++++++++++++++
 2 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/agent/consul/catalog_endpoint_test.go b/agent/consul/catalog_endpoint_test.go
index bd9e125ee844..a44a9e477643 100644
--- a/agent/consul/catalog_endpoint_test.go
+++ b/agent/consul/catalog_endpoint_test.go
@@ -1514,20 +1514,17 @@ func TestCatalog_ListServices_Stale(t *testing.T) {
 		t.Fatalf("bad: %#v", out.Services)
 	}
 
-	if out.Services["consul"] == nil {
-		t.Fatalf("expected consul service.ID, had: %#v", out)
-	}
-
 	if !out.KnownLeader {
 		t.Fatalf("should have a leader: %v", out)
 	}
 
 	s1.Leave()
 	s1.Shutdown()
-	os.RemoveAll(dir1)
+
+	testrpc.WaitUntilNoLeader(t, s2.RPC, "dc1")
 
 	args.AllowStale = false
-	// Run the query, do not wait for leader, never any contact with leader, should fail
+	// Since the leader is now down, non-stale query should fail now
 	if err := msgpackrpc.CallWithCodec(codec, "Catalog.ListServices", &args, &out); err == nil || err.Error() != structs.ErrNoLeader.Error() {
 		t.Fatalf("expected %v but got err: %v and %v", structs.ErrNoLeader, err, out)
 	}
@@ -1543,8 +1540,8 @@ func TestCatalog_ListServices_Stale(t *testing.T) {
 		t.Fatalf("bad: %#v", out)
 	}
 
-	if out.Services["consul"] == nil {
-		t.Fatalf("expected consul service.ID, had: %#v", out)
+	if out.KnownLeader {
+		t.Fatalf("should not have a leader anymore: %#v", out)
 	}
 }
 
diff --git a/testrpc/wait.go b/testrpc/wait.go
index 90d7481bc7c5..b0e6e2ec0e89 100644
--- a/testrpc/wait.go
+++ b/testrpc/wait.go
@@ -26,6 +26,20 @@ func WaitForLeader(t *testing.T, rpc rpcFn, dc string) {
 	})
 }
 
+// WaitUntilNoLeader ensures no leader is present, useful for testing lost leadership.
+func WaitUntilNoLeader(t *testing.T, rpc rpcFn, dc string) {
+	var out structs.IndexedNodes
+	retry.Run(t, func(r *retry.R) {
+		args := &structs.DCSpecificRequest{Datacenter: dc}
+		if err := rpc("Catalog.ListNodes", args, &out); err == nil {
+			r.Fatalf("It still has a leader: %#q", out)
+		}
+		if out.QueryMeta.KnownLeader {
+			r.Fatalf("Has still a leader")
+		}
+	})
+}
+
 // WaitForTestAgent ensures we have a node with serfHealth check registered
 func WaitForTestAgent(t *testing.T, rpc rpcFn, dc string) {
 	var nodes structs.IndexedNodes