From 8b64574721f3dbcd8a4ce4de1cb74104e8704e2a Mon Sep 17 00:00:00 2001
From: Lang Martin <lang@hashicorp.com>
Date: Mon, 20 May 2019 10:54:28 -0400
Subject: [PATCH] e2e util split new alloc and await placement, new
 WaitForDeployment

---
 e2e/e2eutil/utils.go | 39 +++++++++++++++++++++++++++++++++------
 1 file changed, 33 insertions(+), 6 deletions(-)

diff --git a/e2e/e2eutil/utils.go b/e2e/e2eutil/utils.go
index ad17bfa68..c35e52415 100644
--- a/e2e/e2eutil/utils.go
+++ b/e2e/e2eutil/utils.go
@@ -54,15 +54,13 @@ func WaitForNodesReady(t *testing.T, nomadClient *api.Client, nodes int) {
 	})
 }
 
-func RegisterAndWaitForAllocs(t *testing.T, nomadClient *api.Client, jobFile string, jobID string) []*api.AllocationListStub {
+func RegisterAllocs(t *testing.T, nomadClient *api.Client, jobFile string, jobID string) []*api.AllocationListStub {
 	// Parse job
 	job, err := jobspec.ParseFile(jobFile)
 	require := require.New(t)
 	require.Nil(err)
 	job.ID = helper.StringToPtr(jobID)
 
-	g := NewGomegaWithT(t)
-
 	// Register job
 	jobs := nomadClient.Jobs()
 	testutil.WaitForResult(func() (bool, error) {
@@ -75,15 +73,26 @@ func RegisterAndWaitForAllocs(t *testing.T, nomadClient *api.Client, jobFile str
 		require.NoError(err)
 	})
 
+	allocs, _, _ := jobs.Allocations(jobID, false, nil)
+	return allocs
+}
+
+func RegisterAndWaitForAllocs(t *testing.T, nomadClient *api.Client, jobFile string, jobID string) []*api.AllocationListStub {
+	// Parse job
+	require := require.New(t)
+	allocs := RegisterAllocs(t, nomadClient, jobFile, jobID)
+	g := NewGomegaWithT(t)
+	jobs := nomadClient.Jobs()
+
 	// Wrap in retry to wait until placement
 	g.Eventually(func() []*api.AllocationListStub {
 		// Look for allocations
-		allocs, _, _ := jobs.Allocations(*job.ID, false, nil)
+		allocs, _, _ := jobs.Allocations(jobID, false, nil)
 		return allocs
 	}, 30*time.Second, time.Second).ShouldNot(BeEmpty())
 
-	allocs, _, err := jobs.Allocations(*job.ID, false, nil)
-	require.Nil(err)
+	allocs, _, err := jobs.Allocations(jobID, false, nil)
+	require.NoError(err)
 	return allocs
 }
 
@@ -100,3 +109,21 @@ func WaitForAllocRunning(t *testing.T, nomadClient *api.Client, allocID string)
 		t.Fatalf("failed to wait on alloc: %v", err)
 	})
 }
+
+func WaitForDeployment(t *testing.T, nomadClient *api.Client, deployID string, status string, statusDesc string) {
+	testutil.WaitForResultRetries(retries, func() (bool, error) {
+		time.Sleep(time.Millisecond * 100)
+		deploy, _, err := nomadClient.Deployments().Info(deployID, nil)
+		if err != nil {
+			return false, err
+		}
+
+		if deploy.Status == status && deploy.StatusDescription == statusDesc {
+			return true, nil
+		}
+		return false, fmt.Errorf("expected status %s \"%s\", but got: %s \"%s\"", deploy.Status, deploy.StatusDescription)
+
+	}, func(err error) {
+		t.Fatalf("failed to wait on alloc: %v", err)
+	})
+}