From a9c995b0c74f2386b44e68bd3ee00b62f02ecb4e Mon Sep 17 00:00:00 2001
From: Diptanu Choudhury <diptanuc@gmail.com>
Date: Mon, 25 Jul 2016 14:56:38 -0700
Subject: [PATCH] Added a test for adjustQueuedAllocations

---
 scheduler/generic_sched_test.go |  3 ++-
 scheduler/util_test.go          | 34 +++++++++++++++++++++++++++++++++
 2 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/scheduler/generic_sched_test.go b/scheduler/generic_sched_test.go
index 2f437e3a9..da199aeba 100644
--- a/scheduler/generic_sched_test.go
+++ b/scheduler/generic_sched_test.go
@@ -502,7 +502,8 @@ func TestServiceSched_Plan_Partial_Progress(t *testing.T) {
 	node := mock.Node()
 	noErr(t, h.State.UpsertNode(h.NextIndex(), node))
 
-	// Create a job
+	// Create a job with a high resource ask so that all the allocations can't
+	// be placed on a single node.
 	job := mock.Job()
 	job.TaskGroups[0].Count = 3
 	job.TaskGroups[0].Tasks[0].Resources.CPU = 3600
diff --git a/scheduler/util_test.go b/scheduler/util_test.go
index 7484f301f..adbf6db49 100644
--- a/scheduler/util_test.go
+++ b/scheduler/util_test.go
@@ -913,3 +913,37 @@ func TestDesiredUpdates(t *testing.T) {
 		t.Fatalf("desiredUpdates() returned %#v; want %#v", desired, expected)
 	}
 }
+
+func TestUtil_AdjustQueuedAllocations(t *testing.T) {
+	logger := log.New(os.Stderr, "", log.LstdFlags)
+	alloc1 := mock.Alloc()
+	alloc2 := mock.Alloc()
+	alloc2.CreateIndex = 4
+	alloc3 := mock.Alloc()
+	alloc3.CreateIndex = 3
+	alloc4 := mock.Alloc()
+	alloc4.CreateIndex = 6
+
+	planResult := structs.PlanResult{
+		NodeUpdate: map[string][]*structs.Allocation{
+			"node-1": []*structs.Allocation{alloc1},
+		},
+		NodeAllocation: map[string][]*structs.Allocation{
+			"node-1": []*structs.Allocation{
+				alloc2,
+			},
+			"node-2": []*structs.Allocation{
+				alloc3, alloc4,
+			},
+		},
+		RefreshIndex: 3,
+		AllocIndex:   4,
+	}
+
+	queuedAllocs := map[string]int{"web": 2}
+	adjustQueuedAllocations(logger, &planResult, queuedAllocs)
+
+	if queuedAllocs["web"] != 1 {
+		t.Fatalf("expected: %v, actual: %v", 1, queuedAllocs["web"])
+	}
+}