From f86c6976ed7edfee93dc249f141fdce006f7b1c1 Mon Sep 17 00:00:00 2001
From: Preetha Appan <preetha@hashicorp.com>
Date: Wed, 19 Dec 2018 09:05:01 -0600
Subject: [PATCH] minor refactoring of affinities tests

---
 e2e/affinities/affinities.go                  | 86 +++++--------------
 .../{aff3.nomad => anti_affinities.nomad}     |  0
 .../{aff2.nomad => multiple_affinities.nomad} |  0
 .../{aff1.nomad => single_affinity.nomad}     |  0
 4 files changed, 20 insertions(+), 66 deletions(-)
 rename e2e/affinities/input/{aff3.nomad => anti_affinities.nomad} (100%)
 rename e2e/affinities/input/{aff2.nomad => multiple_affinities.nomad} (100%)
 rename e2e/affinities/input/{aff1.nomad => single_affinity.nomad} (100%)

diff --git a/e2e/affinities/affinities.go b/e2e/affinities/affinities.go
index 75f43f3bb..ddee5130d 100644
--- a/e2e/affinities/affinities.go
+++ b/e2e/affinities/affinities.go
@@ -28,15 +28,14 @@ func init() {
 	})
 }
 
-func (tc *BasicAffinityTest) TestSingleAffinities(f *framework.F) {
+func (tc *BasicAffinityTest) registerAndWaitForAllocs(f *framework.F, jobFile string, prefix string) []*api.AllocationListStub {
 	nomadClient := tc.Nomad()
-
 	// Parse job
-	job, err := jobspec.ParseFile("affinities/input/aff1.nomad")
+	job, err := jobspec.ParseFile(jobFile)
 	require := require.New(f.T())
 	require.Nil(err)
 	uuid := uuid.Generate()
-	jobId := helper.StringToPtr("aff" + uuid[0:8])
+	jobId := helper.StringToPtr(prefix + uuid[0:8])
 	job.ID = jobId
 
 	tc.jobIds = append(tc.jobIds, *jobId)
@@ -54,19 +53,24 @@ func (tc *BasicAffinityTest) TestSingleAffinities(f *framework.F) {
 		// Look for allocations
 		allocs, _, _ := jobs.Allocations(*job.ID, false, nil)
 		return allocs
-	}, 5*time.Second, time.Second).ShouldNot(BeEmpty())
+	}, 10*time.Second, time.Second).ShouldNot(BeEmpty())
 
+	allocs, _, err := jobs.Allocations(*job.ID, false, nil)
+	require.Nil(err)
+	return allocs
+}
+
+func (tc *BasicAffinityTest) TestSingleAffinities(f *framework.F) {
+	allocs := tc.registerAndWaitForAllocs(f, "affinities/input/single_affinity.nomad", "aff")
+
+	nomadClient := tc.Nomad()
 	jobAllocs := nomadClient.Allocations()
-
-	allocs, _, _ := jobs.Allocations(*job.ID, false, nil)
-
+	require := require.New(f.T())
 	// Verify affinity score metadata
 	for _, allocStub := range allocs {
 		alloc, _, err := jobAllocs.Info(allocStub.ID, nil)
 		require.Nil(err)
 		require.NotEmpty(alloc.Metrics.ScoreMetaData)
-		// Expect node affinity score to be 1.0 (dc1 nodes) if found
-		// TODO(preetha) make this more realistic
 		for _, sm := range alloc.Metrics.ScoreMetaData {
 			score, ok := sm.Scores["node-affinity"]
 			if ok {
@@ -78,36 +82,11 @@ func (tc *BasicAffinityTest) TestSingleAffinities(f *framework.F) {
 }
 
 func (tc *BasicAffinityTest) TestMultipleAffinities(f *framework.F) {
+	allocs := tc.registerAndWaitForAllocs(f, "affinities/input/multiple_affinities.nomad", "aff")
+
 	nomadClient := tc.Nomad()
-
-	// Parse job
-	job, err := jobspec.ParseFile("affinities/input/aff2.nomad")
-	require := require.New(f.T())
-	require.Nil(err)
-	uuid := uuid.Generate()
-	jobId := helper.StringToPtr("aff" + uuid[0:8])
-	job.ID = jobId
-
-	tc.jobIds = append(tc.jobIds, *jobId)
-
-	// Register job
-	jobs := nomadClient.Jobs()
-	resp, _, err := jobs.Register(job, nil)
-	require.Nil(err)
-	require.NotEmpty(resp.EvalID)
-
-	g := NewGomegaWithT(f.T())
-
-	// Wrap in retry to wait until placement
-	g.Eventually(func() []*api.AllocationListStub {
-		// Look for allocations
-		allocs, _, _ := jobs.Allocations(*job.ID, false, nil)
-		return allocs
-	}, 5*time.Second, time.Second).ShouldNot(BeEmpty())
-
 	jobAllocs := nomadClient.Allocations()
-
-	allocs, _, _ := jobs.Allocations(*job.ID, false, nil)
+	require := require.New(f.T())
 
 	// Verify affinity score metadata
 	for _, allocStub := range allocs {
@@ -142,36 +121,11 @@ func (tc *BasicAffinityTest) TestMultipleAffinities(f *framework.F) {
 }
 
 func (tc *BasicAffinityTest) TestAntiAffinities(f *framework.F) {
+	allocs := tc.registerAndWaitForAllocs(f, "affinities/input/anti_affinities.nomad", "aff")
+
 	nomadClient := tc.Nomad()
-
-	// Parse job
-	job, err := jobspec.ParseFile("affinities/input/aff3.nomad")
-	require := require.New(f.T())
-	require.Nil(err)
-	uuid := uuid.Generate()
-	jobId := helper.StringToPtr("aff" + uuid[0:8])
-	job.ID = jobId
-
-	tc.jobIds = append(tc.jobIds, *jobId)
-
-	// Register job
-	jobs := nomadClient.Jobs()
-	resp, _, err := jobs.Register(job, nil)
-	require.Nil(err)
-	require.NotEmpty(resp.EvalID)
-
-	g := NewGomegaWithT(f.T())
-
-	// Wrap in retry to wait until placement
-	g.Eventually(func() []*api.AllocationListStub {
-		// Look for allocations
-		allocs, _, _ := jobs.Allocations(*job.ID, false, nil)
-		return allocs
-	}, 5*time.Second, time.Second).ShouldNot(BeEmpty())
-
 	jobAllocs := nomadClient.Allocations()
-
-	allocs, _, _ := jobs.Allocations(*job.ID, false, nil)
+	require := require.New(f.T())
 
 	// Verify affinity score metadata
 	for _, allocStub := range allocs {
diff --git a/e2e/affinities/input/aff3.nomad b/e2e/affinities/input/anti_affinities.nomad
similarity index 100%
rename from e2e/affinities/input/aff3.nomad
rename to e2e/affinities/input/anti_affinities.nomad
diff --git a/e2e/affinities/input/aff2.nomad b/e2e/affinities/input/multiple_affinities.nomad
similarity index 100%
rename from e2e/affinities/input/aff2.nomad
rename to e2e/affinities/input/multiple_affinities.nomad
diff --git a/e2e/affinities/input/aff1.nomad b/e2e/affinities/input/single_affinity.nomad
similarity index 100%
rename from e2e/affinities/input/aff1.nomad
rename to e2e/affinities/input/single_affinity.nomad