update task hook coordinator tests

client/allocrunner/task_hook_coordinator_test.go

@@ -2,6 +2,9 @@ package allocrunner
 
 import (
 	"testing"
+	"time"
+
+	"github.com/hashicorp/nomad/nomad/structs"
 
 	"github.com/hashicorp/nomad/helper/testlog"
 	"github.com/hashicorp/nomad/nomad/mock"
@@ -17,41 +20,112 @@ func TestTaskHookCoordinator_OnlyMainApp(t *testing.T) {
 
 	ch := coord.startConditionForTask(tasks[0])
 
-	select {
-	case _, ok := <-ch:
-		require.False(t, ok)
-	default:
-		require.Fail(t, "channel wasn't closed")
-	}
+	testChannelClosed(t, ch, tasks[0].Name)
 }
 
-func TestTaskHookCoordinator_Prestart(t *testing.T) {
-	alloc := mock.Alloc()
-	tasks := alloc.Job.TaskGroups[0].Tasks
+func TestTaskHookCoordinator_PrestartRunsBeforeMain(t *testing.T) {
 	logger := testlog.HCLogger(t)
 
-	tasks = append(tasks, mock.InitTask())
-	tasks = append(tasks, mock.SidecarTask())
-	coord := newTaskHookCoordinator(logger, tasks)
+	alloc := mock.Alloc()
+	tasks := alloc.Job.TaskGroups[0].Tasks
+	tasks = append(tasks, initTask())
+	tasks = append(tasks, sidecarTask())
 
+	coord := newTaskHookCoordinator(logger, tasks)
 	mainCh := coord.startConditionForTask(tasks[0])
 	initCh := coord.startConditionForTask(tasks[1])
 	sideCh := coord.startConditionForTask(tasks[2])
 
-	select {
-	case _, ok := <-initCh:
-		require.False(t, ok)
-	case _, ok := <-sideCh:
-		require.False(t, ok)
-	default:
-		require.Fail(t, "prestart channels weren't closed")
+	mainTaskName := tasks[0].Name
+	initTaskName := tasks[1].Name
+	sideTaskName := tasks[2].Name
+
+	testChannelOpen(t, mainCh, mainTaskName)
+	testChannelClosed(t, initCh, initTaskName)
+	testChannelClosed(t, sideCh, sideTaskName)
+}
+
+func TestTaskHookCoordinator_MainRunsAfterPrestart(t *testing.T) {
+	logger := testlog.HCLogger(t)
+
+	alloc := mock.Alloc()
+	tasks := alloc.Job.TaskGroups[0].Tasks
+	tasks = append(tasks, initTask())
+	tasks = append(tasks, sidecarTask())
+
+	coord := newTaskHookCoordinator(logger, tasks)
+	mainCh := coord.startConditionForTask(tasks[0])
+	initCh := coord.startConditionForTask(tasks[1])
+	sideCh := coord.startConditionForTask(tasks[2])
+
+	mainTaskName := tasks[0].Name
+	initTaskName := tasks[1].Name
+	sideTaskName := tasks[2].Name
+
+	testChannelOpen(t, mainCh, mainTaskName)
+	testChannelClosed(t, initCh, initTaskName)
+	testChannelClosed(t, sideCh, sideTaskName)
+
+	states := map[string]*structs.TaskState{
+		mainTaskName: &structs.TaskState{
+			State:  structs.TaskStatePending,
+			Failed: false,
+		},
+		initTaskName: &structs.TaskState{
+			State:      structs.TaskStateDead,
+			Failed:     false,
+			StartedAt:  time.Now(),
+			FinishedAt: time.Now(),
+		},
+		sideTaskName: &structs.TaskState{
+			State:     structs.TaskStateRunning,
+			Failed:    false,
+			StartedAt: time.Now(),
+		},
 	}
 
+	coord.taskStateUpdated(states)
+
+	testChannelClosed(t, mainCh, mainTaskName)
+	testChannelClosed(t, initCh, initTaskName)
+	testChannelClosed(t, sideCh, sideTaskName)
+}
+
+func testChannelOpen(t *testing.T, ch <-chan struct{}, name string) {
 	select {
-	case <-mainCh:
-		require.Fail(t, "channel was closed, should be open")
+	case <-ch:
+		require.Failf(t, "channel was closed, should be open", name)
 	default:
-		// channel for main task is open, which is correct: coordinator should
-		// block all other tasks until prestart tasks are completed
+		// channel is open
+	}
+}
+
+func testChannelClosed(t *testing.T, ch <-chan struct{}, name string) {
+	select {
+	case _, ok := <-ch:
+		require.False(t, ok)
+	default:
+		// channel is open
+		require.Failf(t, "channel was open, should be closed", name)
+	}
+}
+
+func sidecarTask() *structs.Task {
+	return &structs.Task{
+		Name: "sidecar",
+		Lifecycle: &structs.TaskLifecycleConfig{
+			Hook:       structs.TaskLifecycleHookPrestart,
+			BlockUntil: structs.TaskLifecycleBlockUntilRunning,
+		},
+	}
+}
+
+func initTask() *structs.Task {
+	return &structs.Task{
+		Name: "init",
+		Lifecycle: &structs.TaskLifecycleConfig{
+			Hook:       structs.TaskLifecycleHookPrestart,
+			BlockUntil: structs.TaskLifecycleBlockUntilCompleted,
+		},
 	}
 }

nomad/mock/mock.go

@@ -1013,21 +1013,3 @@ func ACLManagementToken() *structs.ACLToken {
 		ModifyIndex: 20,
 	}
 }
-
-func SidecarTask() *structs.Task {
-	return &structs.Task{
-		Lifecycle: &structs.TaskLifecycleConfig{
-			Hook:       structs.TaskLifecycleHookPrestart,
-			BlockUntil: structs.TaskLifecycleBlockUntilRunning,
-		},
-	}
-}
-
-func InitTask() *structs.Task {
-	return &structs.Task{
-		Lifecycle: &structs.TaskLifecycleConfig{
-			Hook:       structs.TaskLifecycleHookPrestart,
-			BlockUntil: structs.TaskLifecycleBlockUntilCompleted,
-		},
-	}
-}