From eb0840546734475cf02d00e8434f06c7b648abde Mon Sep 17 00:00:00 2001
From: Diptanu Choudhury <diptanuc@gmail.com>
Date: Thu, 28 Jul 2016 14:02:50 -0700
Subject: [PATCH] Updated tests and added logic to system sched

---
 scheduler/generic_sched_test.go |  2 +-
 scheduler/system_sched.go       |  5 ++++
 scheduler/system_sched_test.go  | 41 +++++++++++++++++++++++++++++++++
 3 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/scheduler/generic_sched_test.go b/scheduler/generic_sched_test.go
index 69e72e69d..df1c63a16 100644
--- a/scheduler/generic_sched_test.go
+++ b/scheduler/generic_sched_test.go
@@ -1337,7 +1337,7 @@ func TestServiceSched_NodeUpdate(t *testing.T) {
 	if err != nil {
 		t.Fatalf("err: %v", err)
 	}
-	if h.Evals[0].QueuedAllocations["web"] != 0 {
+	if val, ok := h.Evals[0].QueuedAllocations["web"]; !ok || val != 0 {
 		t.Fatalf("bad queued allocations: %v", h.Evals[0].QueuedAllocations)
 	}
 
diff --git a/scheduler/system_sched.go b/scheduler/system_sched.go
index 812a8b16e..7994efff7 100644
--- a/scheduler/system_sched.go
+++ b/scheduler/system_sched.go
@@ -224,6 +224,11 @@ func (s *SystemScheduler) computeJobAllocs() error {
 
 	// Nothing remaining to do if placement is not required
 	if len(diff.place) == 0 {
+		if s.job != nil {
+			for _, tg := range s.job.TaskGroups {
+				s.queuedAllocs[tg.Name] = 0
+			}
+		}
 		return nil
 	}
 
diff --git a/scheduler/system_sched_test.go b/scheduler/system_sched_test.go
index 46865cccf..f273e0c7d 100644
--- a/scheduler/system_sched_test.go
+++ b/scheduler/system_sched_test.go
@@ -756,6 +756,47 @@ func TestSystemSched_NodeDrain(t *testing.T) {
 	h.AssertEvalStatus(t, structs.EvalStatusComplete)
 }
 
+func TestSystemSched_NodeUpdate(t *testing.T) {
+	h := NewHarness(t)
+
+	// Register a node
+	node := mock.Node()
+	noErr(t, h.State.UpsertNode(h.NextIndex(), node))
+
+	// Generate a fake job allocated on that node.
+	job := mock.SystemJob()
+	noErr(t, h.State.UpsertJob(h.NextIndex(), job))
+
+	alloc := mock.Alloc()
+	alloc.Job = job
+	alloc.JobID = job.ID
+	alloc.NodeID = node.ID
+	alloc.Name = "my-job.web[0]"
+	noErr(t, h.State.UpsertAllocs(h.NextIndex(), []*structs.Allocation{alloc}))
+
+	// Create a mock evaluation to deal
+	eval := &structs.Evaluation{
+		ID:          structs.GenerateUUID(),
+		Priority:    50,
+		TriggeredBy: structs.EvalTriggerNodeUpdate,
+		JobID:       job.ID,
+		NodeID:      node.ID,
+	}
+
+	// Process the evaluation
+	err := h.Process(NewSystemScheduler, eval)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	// Ensure that queued allocations is zero
+	if val, ok := h.Evals[0].QueuedAllocations["web"]; !ok || val != 0 {
+		t.Fatalf("bad queued allocations: %#v", h.Evals[0].QueuedAllocations)
+	}
+
+	h.AssertEvalStatus(t, structs.EvalStatusComplete)
+}
+
 func TestSystemSched_RetryLimit(t *testing.T) {
 	h := NewHarness(t)
 	h.Planner = &RejectPlan{h}