From 9890ac758377d470b4f8c4427ec94a2c4acdae84 Mon Sep 17 00:00:00 2001
From: Preetha Appan <preetha@hashicorp.com>
Date: Wed, 19 Dec 2018 09:56:44 -0600
Subject: [PATCH] Make affinity e2e tests wait for leader through Before method

---
 e2e/affinities/affinities.go |  5 +++++
 e2e/framework/utils.go       | 21 +++++++++++++++++++++
 2 files changed, 26 insertions(+)
 create mode 100644 e2e/framework/utils.go

diff --git a/e2e/affinities/affinities.go b/e2e/affinities/affinities.go
index ddee5130d..f84ab41ce 100644
--- a/e2e/affinities/affinities.go
+++ b/e2e/affinities/affinities.go
@@ -28,6 +28,11 @@ func init() {
 	})
 }
 
+func (tc *BasicAffinityTest) BeforeAll(f *framework.F) {
+	// Ensure cluster has leader before running tests
+	framework.WaitForLeader(f.T(), tc.Nomad())
+}
+
 func (tc *BasicAffinityTest) registerAndWaitForAllocs(f *framework.F, jobFile string, prefix string) []*api.AllocationListStub {
 	nomadClient := tc.Nomad()
 	// Parse job
diff --git a/e2e/framework/utils.go b/e2e/framework/utils.go
new file mode 100644
index 000000000..43ac69c4e
--- /dev/null
+++ b/e2e/framework/utils.go
@@ -0,0 +1,21 @@
+package framework
+
+import (
+	"github.com/hashicorp/nomad/api"
+	"github.com/hashicorp/nomad/testutil"
+	"testing"
+)
+
+// retries is used to control how many times to retry checking if the cluster has a leader yet
+const retries = 500
+
+func WaitForLeader(t *testing.T, nomadClient *api.Client) {
+	statusAPI := nomadClient.Status()
+
+	testutil.WaitForResultRetries(retries, func() (bool, error) {
+		leader, err := statusAPI.Leader()
+		return leader != "", err
+	}, func(err error) {
+		t.Fatalf("failed to find leader: %v", err)
+	})
+}