From 3e8117e91c8df97fd7ff3856303052fabf37f65d Mon Sep 17 00:00:00 2001
From: Armon Dadgar <armon.dadgar@gmail.com>
Date: Thu, 13 Aug 2015 11:33:58 -0700
Subject: [PATCH] scheduler: Adding State to Context

---
 scheduler/context.go       | 13 ++++++++++++-
 scheduler/context_test.go  | 18 ++++++++++++++++++
 scheduler/feasible_test.go |  6 +++---
 scheduler/rank_test.go     |  5 ++++-
 scheduler/select_test.go   |  4 ++--
 5 files changed, 39 insertions(+), 7 deletions(-)
 create mode 100644 scheduler/context_test.go

diff --git a/scheduler/context.go b/scheduler/context.go
index 7a665e2f0..ff661f18a 100644
--- a/scheduler/context.go
+++ b/scheduler/context.go
@@ -2,14 +2,25 @@ package scheduler
 
 // Context is used to track contextual information used for placement
 type Context interface {
+	// State is used to inspect the current global state
+	State() State
 }
 
 // EvalContext is a Context used during an Evaluation
 type EvalContext struct {
+	state State
 }
 
 // NewEvalContext constructs a new EvalContext
-func NewEvalContext() *EvalContext {
+func NewEvalContext(s State) *EvalContext {
 	ctx := &EvalContext{}
 	return ctx
 }
+
+func (e *EvalContext) State() State {
+	return e.state
+}
+
+func (e *EvalContext) SetState(s State) {
+	e.state = s
+}
diff --git a/scheduler/context_test.go b/scheduler/context_test.go
new file mode 100644
index 000000000..c3efa89b1
--- /dev/null
+++ b/scheduler/context_test.go
@@ -0,0 +1,18 @@
+package scheduler
+
+import (
+	"os"
+	"testing"
+
+	"github.com/hashicorp/nomad/nomad/state"
+)
+
+func testContext(t *testing.T) (*state.StateStore, *EvalContext) {
+	state, err := state.NewStateStore(os.Stderr)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	ctx := NewEvalContext(state)
+	return state, ctx
+}
diff --git a/scheduler/feasible_test.go b/scheduler/feasible_test.go
index 841accc26..fcaf6bbd0 100644
--- a/scheduler/feasible_test.go
+++ b/scheduler/feasible_test.go
@@ -9,7 +9,7 @@ import (
 )
 
 func TestRandomIterator(t *testing.T) {
-	ctx := NewEvalContext()
+	_, ctx := testContext(t)
 	var nodes []*structs.Node
 	for i := 0; i < 10; i++ {
 		nodes = append(nodes, mock.Node())
@@ -37,7 +37,7 @@ func TestRandomIterator(t *testing.T) {
 }
 
 func TestDriverIterator(t *testing.T) {
-	ctx := NewEvalContext()
+	_, ctx := testContext(t)
 	nodes := []*structs.Node{
 		mock.Node(),
 		mock.Node(),
@@ -72,7 +72,7 @@ func TestDriverIterator(t *testing.T) {
 }
 
 func TestConstraintIterator(t *testing.T) {
-	ctx := NewEvalContext()
+	_, ctx := testContext(t)
 	nodes := []*structs.Node{
 		mock.Node(),
 		mock.Node(),
diff --git a/scheduler/rank_test.go b/scheduler/rank_test.go
index b2c7775e3..7f3458c53 100644
--- a/scheduler/rank_test.go
+++ b/scheduler/rank_test.go
@@ -8,7 +8,7 @@ import (
 )
 
 func TestFeasibleRankIterator(t *testing.T) {
-	ctx := NewEvalContext()
+	_, ctx := testContext(t)
 	var nodes []*structs.Node
 	for i := 0; i < 10; i++ {
 		nodes = append(nodes, mock.Node())
@@ -31,6 +31,9 @@ func TestFeasibleRankIterator(t *testing.T) {
 	}
 }
 
+func TestBinPackIterator(t *testing.T) {
+}
+
 func TestScoreFit(t *testing.T) {
 	node := mock.Node()
 	node.Resources = &structs.Resources{
diff --git a/scheduler/select_test.go b/scheduler/select_test.go
index 591ebf8f6..d04eb805a 100644
--- a/scheduler/select_test.go
+++ b/scheduler/select_test.go
@@ -7,7 +7,7 @@ import (
 )
 
 func TestLimitIterator(t *testing.T) {
-	ctx := NewEvalContext()
+	_, ctx := testContext(t)
 	nodes := []*RankedNode{
 		&RankedNode{
 			Node:  mock.Node(),
@@ -44,7 +44,7 @@ func TestLimitIterator(t *testing.T) {
 }
 
 func TestMaxScoreIterator(t *testing.T) {
-	ctx := NewEvalContext()
+	_, ctx := testContext(t)
 	nodes := []*RankedNode{
 		&RankedNode{
 			Node:  mock.Node(),