Fix test

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)

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{

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)
+	}
+}