From 5574b2f3d09fd8998fb3c6befc302cd192d50316 Mon Sep 17 00:00:00 2001
From: Mahmood Ali <mahmood@hashicorp.com>
Date: Thu, 6 Jun 2019 15:50:23 -0400
Subject: [PATCH] tests: Migrated allocs aren't lost

Fix `TestServiceSched_NodeDown` for checking that the migrated allocs
are actually marked to be stopped.

The boolean logic in test made it skip actually checking client status
as long as desired status was stop.

Here, we mark some jobs for migration while leaving others as running,
and we check that lost flag is only set for non-migrated allocs.
---
 scheduler/generic_sched_test.go | 36 +++++++++++++++++++++------------
 1 file changed, 23 insertions(+), 13 deletions(-)

diff --git a/scheduler/generic_sched_test.go b/scheduler/generic_sched_test.go
index 70e54d237..8a9acc9bb 100644
--- a/scheduler/generic_sched_test.go
+++ b/scheduler/generic_sched_test.go
@@ -2448,9 +2448,16 @@ func TestServiceSched_NodeDown(t *testing.T) {
 	}
 
 	// Mark appropriate allocs for migration
-	for i := 0; i < 7; i++ {
+	toBeMigrated := map[string]bool{}
+	for i := 0; i < 3; i++ {
 		out := allocs[i]
 		out.DesiredTransition.Migrate = helper.BoolToPtr(true)
+		toBeMigrated[out.ID] = true
+	}
+
+	toBeLost := map[string]bool{}
+	for i := len(toBeMigrated); i < 7; i++ {
+		toBeLost[allocs[i].ID] = true
 	}
 
 	noErr(t, h.State.UpsertAllocs(h.NextIndex(), allocs))
@@ -2469,25 +2476,28 @@ func TestServiceSched_NodeDown(t *testing.T) {
 
 	// Process the evaluation
 	err := h.Process(NewServiceScheduler, eval)
-	if err != nil {
-		t.Fatalf("err: %v", err)
-	}
+	require.NoError(t, err)
 
 	// Ensure a single plan
-	if len(h.Plans) != 1 {
-		t.Fatalf("bad: %#v", h.Plans)
-	}
+	require.Len(t, h.Plans, 1)
 	plan := h.Plans[0]
 
 	// Test the scheduler marked all non-terminal allocations as lost
-	if len(plan.NodeUpdate[node.ID]) != 7 {
-		t.Fatalf("bad: %#v", plan)
-	}
+	require.Len(t, plan.NodeUpdate[node.ID], len(toBeMigrated)+len(toBeLost))
 
 	for _, out := range plan.NodeUpdate[node.ID] {
-		if out.ClientStatus != structs.AllocClientStatusLost && out.DesiredStatus != structs.AllocDesiredStatusStop {
-			t.Fatalf("bad alloc: %#v", out)
-		}
+		t.Run("alloc "+out.ID, func(t *testing.T) {
+			require.Equal(t, structs.AllocDesiredStatusStop, out.DesiredStatus)
+
+			if toBeMigrated[out.ID] {
+				// there is no indicator on job itself that marks it as migrated
+				require.NotEqual(t, structs.AllocClientStatusLost, out.ClientStatus)
+			} else if toBeLost[out.ID] {
+				require.Equal(t, structs.AllocClientStatusLost, out.ClientStatus)
+			} else {
+				require.Fail(t, "unexpected alloc update")
+			}
+		})
 	}
 
 	h.AssertEvalStatus(t, structs.EvalStatusComplete)