From a5b024fdea1f8552b1891546575c556160b40ae6 Mon Sep 17 00:00:00 2001
From: Mahmood Ali <mahmood@hashicorp.com>
Date: Tue, 24 Mar 2020 21:52:39 -0400
Subject: [PATCH] tests: restart restartpolicy for all tasks in tests

---
 client/allocrunner/alloc_runner_unix_test.go      |  6 +++---
 client/allocrunner/taskrunner/task_runner_test.go | 10 ++++++++--
 2 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/client/allocrunner/alloc_runner_unix_test.go b/client/allocrunner/alloc_runner_unix_test.go
index 7cc830510..7ec4dea0b 100644
--- a/client/allocrunner/alloc_runner_unix_test.go
+++ b/client/allocrunner/alloc_runner_unix_test.go
@@ -236,9 +236,9 @@ func TestAllocRunner_PreStartFailuresLeadToFailed(t *testing.T) {
 	task.Config = map[string]interface{}{
 		"run_for": "2ms",
 	}
-	alloc.Job.TaskGroups[0].RestartPolicy = &structs.RestartPolicy{
-		Attempts: 0,
-	}
+	rp := &structs.RestartPolicy{Attempts: 0}
+	alloc.Job.TaskGroups[0].RestartPolicy = rp
+	task.RestartPolicy = rp
 
 	conf, cleanup := testAllocRunnerConfig(t, alloc.Copy())
 	defer cleanup()
diff --git a/client/allocrunner/taskrunner/task_runner_test.go b/client/allocrunner/taskrunner/task_runner_test.go
index 0fe7ac112..fd9c4f113 100644
--- a/client/allocrunner/taskrunner/task_runner_test.go
+++ b/client/allocrunner/taskrunner/task_runner_test.go
@@ -1502,6 +1502,7 @@ func TestTaskRunner_Download_ChrootExec(t *testing.T) {
 	alloc := mock.BatchAlloc()
 	alloc.Job.TaskGroups[0].RestartPolicy = &structs.RestartPolicy{}
 	task := alloc.Job.TaskGroups[0].Tasks[0]
+	task.RestartPolicy = &structs.RestartPolicy{}
 	task.Driver = "exec"
 	task.Config = map[string]interface{}{
 		"command": "noop.sh",
@@ -1541,6 +1542,7 @@ func TestTaskRunner_Download_RawExec(t *testing.T) {
 	alloc := mock.BatchAlloc()
 	alloc.Job.TaskGroups[0].RestartPolicy = &structs.RestartPolicy{}
 	task := alloc.Job.TaskGroups[0].Tasks[0]
+	task.RestartPolicy = &structs.RestartPolicy{}
 	task.Driver = "raw_exec"
 	task.Config = map[string]interface{}{
 		"command": "noop.sh",
@@ -1631,12 +1633,14 @@ func TestTaskRunner_Download_Retries(t *testing.T) {
 	task.Artifacts = []*structs.TaskArtifact{&artifact}
 
 	// Make the restart policy retry once
-	alloc.Job.TaskGroups[0].RestartPolicy = &structs.RestartPolicy{
+	rp := &structs.RestartPolicy{
 		Attempts: 1,
 		Interval: 10 * time.Minute,
 		Delay:    1 * time.Second,
 		Mode:     structs.RestartPolicyModeFail,
 	}
+	alloc.Job.TaskGroups[0].RestartPolicy = rp
+	alloc.Job.TaskGroups[0].Tasks[0].RestartPolicy = rp
 
 	tr, _, cleanup := runTestTaskRunner(t, alloc, task.Name)
 	defer cleanup()
@@ -1873,12 +1877,14 @@ func TestTaskRunner_Run_RecoverableStartError(t *testing.T) {
 	}
 
 	// Make the restart policy retry once
-	alloc.Job.TaskGroups[0].RestartPolicy = &structs.RestartPolicy{
+	rp := &structs.RestartPolicy{
 		Attempts: 1,
 		Interval: 10 * time.Minute,
 		Delay:    0,
 		Mode:     structs.RestartPolicyModeFail,
 	}
+	alloc.Job.TaskGroups[0].RestartPolicy = rp
+	alloc.Job.TaskGroups[0].Tasks[0].RestartPolicy = rp
 
 	tr, _, cleanup := runTestTaskRunner(t, alloc, task.Name)
 	defer cleanup()