From 3040c4dab267912f301ba520c154f8300145360e Mon Sep 17 00:00:00 2001
From: Alex Dadgar <alex.dadgar@gmail.com>
Date: Tue, 8 Mar 2016 11:07:16 -0800
Subject: [PATCH] Fix test

---
 client/driver/docker_test.go | 19 ++++++++++---------
 client/driver/driver_test.go | 18 ------------------
 client/driver/utils_test.go  | 22 ++++++++++++++++++++++
 3 files changed, 32 insertions(+), 27 deletions(-)
 create mode 100644 client/driver/utils_test.go

diff --git a/client/driver/docker_test.go b/client/driver/docker_test.go
index b0cbd615216..9e641c219b3 100644
--- a/client/driver/docker_test.go
+++ b/client/driver/docker_test.go
@@ -152,18 +152,19 @@ func TestDockerDriver_Handle(t *testing.T) {
 	defer pluginClient.Kill()
 
 	h := &DockerHandle{
-		version:      "version",
-		imageID:      "imageid",
-		logCollector: logCollector,
-		pluginClient: pluginClient,
-		containerID:  "containerid",
-		killTimeout:  5 * time.Nanosecond,
-		doneCh:       make(chan struct{}),
-		waitCh:       make(chan *cstructs.WaitResult, 1),
+		version:        "version",
+		imageID:        "imageid",
+		logCollector:   logCollector,
+		pluginClient:   pluginClient,
+		containerID:    "containerid",
+		killTimeout:    5 * time.Nanosecond,
+		maxKillTimeout: 15 * time.Nanosecond,
+		doneCh:         make(chan struct{}),
+		waitCh:         make(chan *cstructs.WaitResult, 1),
 	}
 
 	actual := h.ID()
-	expected := fmt.Sprintf("DOCKER:{\"Version\":\"version\",\"ImageID\":\"imageid\",\"ContainerID\":\"containerid\",\"KillTimeout\":5,\"PluginConfig\":{\"Pid\":%d,\"AddrNet\":\"unix\",\"AddrName\":\"%s\"}}",
+	expected := fmt.Sprintf("DOCKER:{\"Version\":\"version\",\"ImageID\":\"imageid\",\"ContainerID\":\"containerid\",\"KillTimeout\":5,\"MaxKillTimeout\":15,\"PluginConfig\":{\"Pid\":%d,\"AddrNet\":\"unix\",\"AddrName\":\"%s\"}}",
 		pluginClient.ReattachConfig().Pid, pluginClient.ReattachConfig().Addr.String())
 	if actual != expected {
 		t.Errorf("Expected `%s`, found `%s`", expected, actual)
diff --git a/client/driver/driver_test.go b/client/driver/driver_test.go
index 89beb40bf28..d1ce7963c0c 100644
--- a/client/driver/driver_test.go
+++ b/client/driver/driver_test.go
@@ -66,24 +66,6 @@ func testDriverContexts(task *structs.Task) (*DriverContext, *ExecContext) {
 	return driverCtx, execCtx
 }
 
-func TestDriver_KillTimeout(t *testing.T) {
-	expected := 1 * time.Second
-	task := &structs.Task{Name: "foo", KillTimeout: expected}
-	ctx, _ := testDriverContexts(task)
-	ctx.config.MaxKillTimeout = 10 * time.Second
-
-	if actual := ctx.KillTimeout(task); expected != actual {
-		t.Fatalf("KillTimeout(%v) returned %v; want %v", task, actual, expected)
-	}
-
-	expected = 10 * time.Second
-	task = &structs.Task{KillTimeout: 11 * time.Second}
-
-	if actual := ctx.KillTimeout(task); expected != actual {
-		t.Fatalf("KillTimeout(%v) returned %v; want %v", task, actual, expected)
-	}
-}
-
 func TestDriver_GetTaskEnv(t *testing.T) {
 	t.Parallel()
 	task := &structs.Task{
diff --git a/client/driver/utils_test.go b/client/driver/utils_test.go
new file mode 100644
index 00000000000..7ea78c680c9
--- /dev/null
+++ b/client/driver/utils_test.go
@@ -0,0 +1,22 @@
+package driver
+
+import (
+	"testing"
+	"time"
+)
+
+func TestDriver_KillTimeout(t *testing.T) {
+	expected := 1 * time.Second
+	max := 10 * time.Second
+
+	if actual := GetKillTimeout(expected, max); expected != actual {
+		t.Fatalf("GetKillTimeout() returned %v; want %v", actual, expected)
+	}
+
+	expected = 10 * time.Second
+	input := 11 * time.Second
+
+	if actual := GetKillTimeout(input, max); expected != actual {
+		t.Fatalf("KillTimeout() returned %v; want %v", actual, expected)
+	}
+}