diff --git a/client/task_runner_test.go b/client/task_runner_test.go
index 52c8c73b3..f0f6c92e2 100644
--- a/client/task_runner_test.go
+++ b/client/task_runner_test.go
@@ -49,7 +49,7 @@ func testTaskRunner(restarts bool) (*MockTaskStateUpdater, *TaskRunner) {
 
 	ctx := driver.NewExecContext(allocDir, alloc.ID)
 	rp := structs.NewRestartPolicy(structs.JobTypeService)
-	restartTracker := newRestartTracker(rp)
+	restartTracker := newRestartTracker(rp, alloc.Job.Type)
 	if !restarts {
 		restartTracker = noRestartsTracker()
 	}