Test task group update annotations

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)

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{}