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

---
 scheduler/context.go       |  4 ++++
 scheduler/feasible.go      |  7 ++-----
 scheduler/feasible_test.go | 37 +++++++++++++++++++++++++++++++++++++
 scheduler/rank.go          |  3 +++
 scheduler/rank_test.go     | 32 ++++++++++++++++++++++++++++++++
 5 files changed, 78 insertions(+), 5 deletions(-)
 create mode 100644 scheduler/feasible_test.go
 create mode 100644 scheduler/rank_test.go

diff --git a/scheduler/context.go b/scheduler/context.go
index 34cfbc446..7a665e2f0 100644
--- a/scheduler/context.go
+++ b/scheduler/context.go
@@ -1,5 +1,9 @@
 package scheduler
 
+// Context is used to track contextual information used for placement
+type Context interface {
+}
+
 // EvalContext is a Context used during an Evaluation
 type EvalContext struct {
 }
diff --git a/scheduler/feasible.go b/scheduler/feasible.go
index e9803a49e..442ce366f 100644
--- a/scheduler/feasible.go
+++ b/scheduler/feasible.go
@@ -6,10 +6,6 @@ import (
 	"github.com/hashicorp/nomad/nomad/structs"
 )
 
-// Context is used to track contextual information used for placement
-type Context interface {
-}
-
 // FeasibleIterator is used to iteratively yield nodes that
 // match feasibility constraints. The iterators may manage
 // some state for performance optimizations.
@@ -49,7 +45,8 @@ func (iter *StaticIterator) Next() *structs.Node {
 }
 
 // NewRandomIterator constructs a static iterator from a list of nodes
-// after applying the Fisher-Yates algorithm for a random shuffle
+// after applying the Fisher-Yates algorithm for a random shuffle. This
+// is applied in-place
 func NewRandomIterator(ctx Context, nodes []*structs.Node) *StaticIterator {
 	// shuffle with the Fisher-Yates algorithm
 	n := len(nodes)
diff --git a/scheduler/feasible_test.go b/scheduler/feasible_test.go
new file mode 100644
index 000000000..bcdc668cc
--- /dev/null
+++ b/scheduler/feasible_test.go
@@ -0,0 +1,37 @@
+package scheduler
+
+import (
+	"reflect"
+	"testing"
+
+	"github.com/hashicorp/nomad/nomad/mock"
+	"github.com/hashicorp/nomad/nomad/structs"
+)
+
+func TestRandomIterator(t *testing.T) {
+	ctx := NewEvalContext()
+	var nodes []*structs.Node
+	for i := 0; i < 10; i++ {
+		nodes = append(nodes, mock.Node())
+	}
+
+	nc := make([]*structs.Node, len(nodes))
+	copy(nc, nodes)
+	rand := NewRandomIterator(ctx, nc)
+
+	var out []*structs.Node
+	for {
+		next := rand.Next()
+		if next == nil {
+			break
+		}
+		out = append(out, next)
+	}
+
+	if len(out) != len(nodes) {
+		t.Fatalf("missing nodes")
+	}
+	if reflect.DeepEqual(out, nodes) {
+		t.Fatalf("same order")
+	}
+}
diff --git a/scheduler/rank.go b/scheduler/rank.go
index d7a121127..70a6208ef 100644
--- a/scheduler/rank.go
+++ b/scheduler/rank.go
@@ -36,6 +36,9 @@ func NewFeasibleRankIterator(ctx Context, source FeasibleIterator) *FeasibleRank
 
 func (iter *FeasibleRankIterator) Next() *RankedNode {
 	option := iter.source.Next()
+	if option == nil {
+		return nil
+	}
 	ranked := &RankedNode{
 		Node: option,
 	}
diff --git a/scheduler/rank_test.go b/scheduler/rank_test.go
new file mode 100644
index 000000000..16a75b992
--- /dev/null
+++ b/scheduler/rank_test.go
@@ -0,0 +1,32 @@
+package scheduler
+
+import (
+	"testing"
+
+	"github.com/hashicorp/nomad/nomad/mock"
+	"github.com/hashicorp/nomad/nomad/structs"
+)
+
+func TestFeasibleRankIterator(t *testing.T) {
+	ctx := NewEvalContext()
+	var nodes []*structs.Node
+	for i := 0; i < 10; i++ {
+		nodes = append(nodes, mock.Node())
+	}
+	static := NewStaticIterator(ctx, nodes)
+
+	feasible := NewFeasibleRankIterator(ctx, static)
+
+	var out []*RankedNode
+	for {
+		next := feasible.Next()
+		if next == nil {
+			break
+		}
+		out = append(out, next)
+	}
+
+	if len(out) != len(nodes) {
+		t.Fatalf("bad: %v", out)
+	}
+}