From 5540b117b82bc19670941808b8034e6b613b33f3 Mon Sep 17 00:00:00 2001
From: Alex Dadgar <alex.dadgar@gmail.com>
Date: Wed, 11 May 2016 16:31:50 -0700
Subject: [PATCH] Test task group update annotations

---
 scheduler/annotate.go      |  2 --
 scheduler/annotate_test.go | 42 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/scheduler/annotate.go b/scheduler/annotate.go
index 559e640b7..461be9906 100644
--- a/scheduler/annotate.go
+++ b/scheduler/annotate.go
@@ -1,7 +1,6 @@
 package scheduler
 
 import (
-	"fmt"
 	"strconv"
 
 	"github.com/hashicorp/nomad/nomad/structs"
@@ -129,7 +128,6 @@ func annotateCountChange(diff *structs.TaskGroupDiff) error {
 
 	// Didn't find
 	if countDiff == nil {
-		fmt.Println("NO COUNT")
 		return nil
 	}
 	oldV, err := strconv.Atoi(countDiff.Old)
diff --git a/scheduler/annotate_test.go b/scheduler/annotate_test.go
index 5f90ef23f..e8713d93e 100644
--- a/scheduler/annotate_test.go
+++ b/scheduler/annotate_test.go
@@ -7,6 +7,48 @@ import (
 	"github.com/hashicorp/nomad/nomad/structs"
 )
 
+func TestAnnotateTaskGroup_Updates(t *testing.T) {
+	plan := &structs.Plan{
+		Annotations: &structs.PlanAnnotations{
+			DesiredTGUpdates: map[string]*structs.DesiredUpdates{
+				"foo": &structs.DesiredUpdates{
+					Ignore:            1,
+					Place:             2,
+					Migrate:           3,
+					Stop:              4,
+					InPlaceUpdate:     5,
+					DestructiveUpdate: 6,
+				},
+			},
+		},
+	}
+
+	tgDiff := &structs.TaskGroupDiff{
+		Type: structs.DiffTypeEdited,
+		Name: "foo",
+	}
+	expected := &structs.TaskGroupDiff{
+		Type: structs.DiffTypeEdited,
+		Name: "foo",
+		Updates: map[string]uint64{
+			UpdateTypeIgnore:            1,
+			UpdateTypeCreate:            2,
+			UpdateTypeMigrate:           3,
+			UpdateTypeDestroy:           4,
+			UpdateTypeInplaceUpdate:     5,
+			UpdateTypeDestructiveUpdate: 6,
+		},
+	}
+
+	if err := annotateTaskGroup(tgDiff, plan); err != nil {
+		t.Fatalf("annotateTaskGroup(%#v, %#v) failed: %#v", tgDiff, plan, err)
+	}
+
+	if !reflect.DeepEqual(tgDiff, expected) {
+		t.Fatalf("got %#v, want %#v", tgDiff, expected)
+	}
+}
+
 func TestAnnotateCountChange_NonEdited(t *testing.T) {
 	tg := &structs.TaskGroupDiff{}
 	tgOrig := &structs.TaskGroupDiff{}