scheduler: testing select iterators

scheduler/context.go

@@ -0,0 +1,11 @@
+package scheduler
+
+// EvalContext is a Context used during an Evaluation
+type EvalContext struct {
+}
+
+// NewEvalContext constructs a new EvalContext
+func NewEvalContext() *EvalContext {
+	ctx := &EvalContext{}
+	return ctx
+}

scheduler/rank.go

@@ -50,6 +50,15 @@ type StaticRankIterator struct {
 	offset int
 }
 
+// NewStaticRankIterator returns a new static rank iterator over the given nodes
+func NewStaticRankIterator(ctx Context, nodes []*RankedNode) *StaticRankIterator {
+	iter := &StaticRankIterator{
+		ctx:   ctx,
+		nodes: nodes,
+	}
+	return iter
+}
+
 func (iter *StaticRankIterator) Next() *RankedNode {
 	// Check if exhausted
 	if iter.offset == len(iter.nodes) {

scheduler/select.go

@@ -42,21 +42,30 @@ type MaxScoreIterator struct {
 	max    *RankedNode
 }
 
+// MaxScoreIterator returns a MaxScoreIterator over the given source
+func NewMaxScoreIterator(ctx Context, source RankIterator) *MaxScoreIterator {
+	iter := &MaxScoreIterator{
+		ctx:    ctx,
+		source: source,
+	}
+	return iter
+}
+
 func (iter *MaxScoreIterator) Next() *RankedNode {
+	// Check if we've found the max, return nil
+	if iter.max != nil {
+		return nil
+	}
+
+	// Consume and determine the max
 	for {
 		option := iter.source.Next()
 		if option == nil {
-			break
+			return iter.max
 		}
 
-		if iter.max == nil {
-			iter.max = option
-			continue
-		}
-
-		if option.Score > iter.max.Score {
+		if iter.max == nil || option.Score > iter.max.Score {
 			iter.max = option
 		}
 	}
-	return iter.max
 }

scheduler/select_test.go

@@ -0,0 +1,81 @@
+package scheduler
+
+import (
+	"testing"
+
+	"github.com/hashicorp/nomad/nomad/mock"
+)
+
+func TestLimitIterator(t *testing.T) {
+	ctx := NewEvalContext()
+	nodes := []*RankedNode{
+		&RankedNode{
+			Node:  mock.Node(),
+			Score: 1,
+		},
+		&RankedNode{
+			Node:  mock.Node(),
+			Score: 2,
+		},
+		&RankedNode{
+			Node:  mock.Node(),
+			Score: 3,
+		},
+	}
+	static := NewStaticRankIterator(ctx, nodes)
+
+	limit := NewLimitIterator(ctx, static, 2)
+
+	var out []*RankedNode
+	for {
+		next := limit.Next()
+		if next == nil {
+			break
+		}
+		out = append(out, next)
+	}
+
+	if len(out) != 2 {
+		t.Fatalf("bad: %v", out)
+	}
+	if out[0] != nodes[0] && out[1] != nodes[1] {
+		t.Fatalf("bad: %v", out)
+	}
+}
+
+func TestMaxScoreIterator(t *testing.T) {
+	ctx := NewEvalContext()
+	nodes := []*RankedNode{
+		&RankedNode{
+			Node:  mock.Node(),
+			Score: 1,
+		},
+		&RankedNode{
+			Node:  mock.Node(),
+			Score: 2,
+		},
+		&RankedNode{
+			Node:  mock.Node(),
+			Score: 3,
+		},
+	}
+	static := NewStaticRankIterator(ctx, nodes)
+
+	max := NewMaxScoreIterator(ctx, static)
+
+	var out []*RankedNode
+	for {
+		next := max.Next()
+		if next == nil {
+			break
+		}
+		out = append(out, next)
+	}
+
+	if len(out) != 1 {
+		t.Fatalf("bad: %v", out)
+	}
+	if out[0] != nodes[2] {
+		t.Fatalf("bad: %v", out)
+	}
+}