From 762db7c5d71f199167fa2f7be04963f525ba8a94 Mon Sep 17 00:00:00 2001
From: Alex Dadgar <alex.dadgar@gmail.com>
Date: Mon, 26 Feb 2018 14:34:32 -0800
Subject: [PATCH] Fix tests

---
 nomad/drain_test.go             |  4 +++-
 nomad/fsm_test.go               |  4 +++-
 nomad/node_endpoint_test.go     | 16 ++++++++++++----
 nomad/state/state_store_test.go |  4 +++-
 4 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/nomad/drain_test.go b/nomad/drain_test.go
index c47e0d401..13465ede0 100644
--- a/nomad/drain_test.go
+++ b/nomad/drain_test.go
@@ -144,7 +144,9 @@ func TestNodeDrainer_SimpleDrain(t *testing.T) {
 	// Start draining node 1
 	//FIXME update drain rpc to skip fsm manipulation and use api
 	strategy := &structs.DrainStrategy{
-		Deadline: -1 * time.Second,
+		DrainSpec: structs.DrainSpec{
+			Deadline: -1 * time.Second,
+		},
 	}
 	node, err := state.NodeByID(nil, c1.NodeID())
 	require.Nil(err)
diff --git a/nomad/fsm_test.go b/nomad/fsm_test.go
index b9fc4845d..87c3a8b44 100644
--- a/nomad/fsm_test.go
+++ b/nomad/fsm_test.go
@@ -294,7 +294,9 @@ func TestFSM_UpdateNodeDrain(t *testing.T) {
 	require.Nil(resp)
 
 	strategy := &structs.DrainStrategy{
-		Deadline: 10 * time.Second,
+		DrainSpec: structs.DrainSpec{
+			Deadline: 10 * time.Second,
+		},
 	}
 	req2 := structs.NodeUpdateDrainRequest{
 		NodeID:        node.ID,
diff --git a/nomad/node_endpoint_test.go b/nomad/node_endpoint_test.go
index 0de46ed22..5ce3cc272 100644
--- a/nomad/node_endpoint_test.go
+++ b/nomad/node_endpoint_test.go
@@ -765,7 +765,9 @@ func TestClientEndpoint_UpdateDrain(t *testing.T) {
 	require.Nil(msgpackrpc.CallWithCodec(codec, "Node.Register", reg, &resp))
 
 	strategy := &structs.DrainStrategy{
-		Deadline: 10 * time.Second,
+		DrainSpec: structs.DrainSpec{
+			Deadline: 10 * time.Second,
+		},
 	}
 
 	// Update the status
@@ -809,7 +811,9 @@ func TestClientEndpoint_UpdateDrain_ACL(t *testing.T) {
 	dereg := &structs.NodeUpdateDrainRequest{
 		NodeID: node.ID,
 		DrainStrategy: &structs.DrainStrategy{
-			Deadline: 10 * time.Second,
+			DrainSpec: structs.DrainSpec{
+				Deadline: 10 * time.Second,
+			},
 		},
 		WriteRequest: structs.WriteRequest{Region: "global"},
 	}
@@ -910,7 +914,9 @@ func TestClientEndpoint_Drain_Down(t *testing.T) {
 	dereg := &structs.NodeUpdateDrainRequest{
 		NodeID: node.ID,
 		DrainStrategy: &structs.DrainStrategy{
-			Deadline: -1 * time.Second,
+			DrainSpec: structs.DrainSpec{
+				Deadline: -1 * time.Second,
+			},
 		},
 		WriteRequest: structs.WriteRequest{Region: "global"},
 	}
@@ -2369,7 +2375,9 @@ func TestClientEndpoint_ListNodes_Blocking(t *testing.T) {
 	// Node drain updates trigger watches.
 	time.AfterFunc(100*time.Millisecond, func() {
 		s := &structs.DrainStrategy{
-			Deadline: 10 * time.Second,
+			DrainSpec: structs.DrainSpec{
+				Deadline: 10 * time.Second,
+			},
 		}
 		if err := state.UpdateNodeDrain(3, node.ID, s, 101); err != nil {
 			t.Fatalf("err: %v", err)
diff --git a/nomad/state/state_store_test.go b/nomad/state/state_store_test.go
index af2e8cceb..3d97cf7f5 100644
--- a/nomad/state/state_store_test.go
+++ b/nomad/state/state_store_test.go
@@ -707,7 +707,9 @@ func TestStateStore_UpdateNodeDrain_Node(t *testing.T) {
 
 	expectedTime := int64(101)
 	expectedDrain := &structs.DrainStrategy{
-		Deadline: 10 * time.Second,
+		DrainSpec: structs.DrainSpec{
+			Deadline: -1 * time.Second,
+		},
 	}
 
 	// Create a watchset so we can test that update node drain fires the watch