From fc48be3656814d58c06a9d85a49295de7d09ad91 Mon Sep 17 00:00:00 2001
From: Preetha Appan <preetha@hashicorp.com>
Date: Sun, 29 Jul 2018 19:49:56 -0500
Subject: [PATCH] added some unit tests for -1 spread score

---
 scheduler/spread.go      |  8 ++--
 scheduler/spread_test.go | 88 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 92 insertions(+), 4 deletions(-)

diff --git a/scheduler/spread.go b/scheduler/spread.go
index e6c3297a9..22c30ccc1 100644
--- a/scheduler/spread.go
+++ b/scheduler/spread.go
@@ -5,9 +5,9 @@ import (
 )
 
 const (
-	// ImplicitTarget is used to represent any remaining attribute values
+	// implicitTarget is used to represent any remaining attribute values
 	// when target percentages don't add up to 100
-	ImplicitTarget = "*"
+	implicitTarget = "*"
 )
 
 // SpreadIterator is used to spread allocations across a specified attribute
@@ -139,7 +139,7 @@ func (iter *SpreadIterator) Next() *RankedNode {
 				desiredCount, ok := spreadDetails.desiredCounts[nValue]
 				if !ok {
 					// See if there is an implicit target
-					desiredCount, ok = spreadDetails.desiredCounts[ImplicitTarget]
+					desiredCount, ok = spreadDetails.desiredCounts[implicitTarget]
 					if !ok {
 						// The desired count for this attribute is zero if it gets here
 						// so use the maximum possible penalty for this node
@@ -241,7 +241,7 @@ func (iter *SpreadIterator) computeSpreadInfo(tg *structs.TaskGroup) {
 		// Account for remaining count only if there is any spread targets
 		if sumDesiredCounts > 0 && sumDesiredCounts < float64(totalCount) {
 			remainingCount := float64(totalCount) - sumDesiredCounts
-			si.desiredCounts[ImplicitTarget] = remainingCount
+			si.desiredCounts[implicitTarget] = remainingCount
 		}
 		spreadInfos[spread.Attribute] = si
 		iter.sumSpreadWeights += spread.Weight
diff --git a/scheduler/spread_test.go b/scheduler/spread_test.go
index 659044b29..f7e7222f7 100644
--- a/scheduler/spread_test.go
+++ b/scheduler/spread_test.go
@@ -431,3 +431,91 @@ func TestSpreadIterator_EvenSpread(t *testing.T) {
 	}
 
 }
+
+// Test scenarios where the spread iterator sets maximum penalty (-1.0)
+func TestSpreadIterator_MaxPenalty(t *testing.T) {
+	state, ctx := testContext(t)
+	var nodes []*RankedNode
+
+	// Add nodes in dc3 to the state store
+	for i := 0; i < 5; i++ {
+		node := mock.Node()
+		node.Datacenter = "dc3"
+		if err := state.UpsertNode(uint64(100+i), node); err != nil {
+			t.Fatalf("failed to upsert node: %v", err)
+		}
+		nodes = append(nodes, &RankedNode{Node: node})
+	}
+
+	static := NewStaticRankIterator(ctx, nodes)
+
+	job := mock.Job()
+	tg := job.TaskGroups[0]
+	job.TaskGroups[0].Count = 5
+
+	// Create spread target of 80% in dc1
+	// and 20% in dc2
+	spread := &structs.Spread{
+		Weight:    100,
+		Attribute: "${node.datacenter}",
+		SpreadTarget: []*structs.SpreadTarget{
+			{
+				Value:   "dc1",
+				Percent: 80,
+			},
+			{
+				Value:   "dc2",
+				Percent: 20,
+			},
+		},
+	}
+	tg.Spreads = []*structs.Spread{spread}
+	spreadIter := NewSpreadIterator(ctx, static)
+	spreadIter.SetJob(job)
+	spreadIter.SetTaskGroup(tg)
+
+	scoreNorm := NewScoreNormalizationIterator(ctx, spreadIter)
+
+	out := collectRanked(scoreNorm)
+
+	// All nodes are in dc3 so score should be -1
+	for _, rn := range out {
+		require.Equal(t, -1.0, rn.FinalScore)
+	}
+
+	// Reset scores
+	for _, node := range nodes {
+		node.Scores = nil
+		node.FinalScore = 0
+	}
+
+	// Create spread on attribute that doesn't exist on any nodes
+	spread = &structs.Spread{
+		Weight:    100,
+		Attribute: "${meta.foo}",
+		SpreadTarget: []*structs.SpreadTarget{
+			{
+				Value:   "bar",
+				Percent: 80,
+			},
+			{
+				Value:   "baz",
+				Percent: 20,
+			},
+		},
+	}
+
+	tg.Spreads = []*structs.Spread{spread}
+	static = NewStaticRankIterator(ctx, nodes)
+	spreadIter = NewSpreadIterator(ctx, static)
+	spreadIter.SetJob(job)
+	spreadIter.SetTaskGroup(tg)
+	scoreNorm = NewScoreNormalizationIterator(ctx, spreadIter)
+	out = collectRanked(scoreNorm)
+
+	// All nodes don't have the spread attribute so score should be -1
+	for _, rn := range out {
+		require.Equal(t, -1.0, rn.FinalScore)
+	}
+
+}