From 6ecf43bdeeaca2a246308ff740ef6adaaeb8a274 Mon Sep 17 00:00:00 2001
From: Armon Dadgar <armon.dadgar@gmail.com>
Date: Thu, 13 Aug 2015 10:05:54 -0700
Subject: [PATCH] scheduler: testing select iterators

---
 scheduler/context.go     | 11 ++++++
 scheduler/rank.go        |  9 +++++
 scheduler/select.go      | 25 +++++++++----
 scheduler/select_test.go | 81 ++++++++++++++++++++++++++++++++++++++++
 4 files changed, 118 insertions(+), 8 deletions(-)
 create mode 100644 scheduler/context.go
 create mode 100644 scheduler/select_test.go

diff --git a/scheduler/context.go b/scheduler/context.go
new file mode 100644
index 000000000..34cfbc446
--- /dev/null
+++ b/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
+}
diff --git a/scheduler/rank.go b/scheduler/rank.go
index 54a82142b..d7a121127 100644
--- a/scheduler/rank.go
+++ b/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) {
diff --git a/scheduler/select.go b/scheduler/select.go
index 97e78edf8..c167282fb 100644
--- a/scheduler/select.go
+++ b/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
 }
diff --git a/scheduler/select_test.go b/scheduler/select_test.go
new file mode 100644
index 000000000..591ebf8f6
--- /dev/null
+++ b/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)
+	}
+}