Repository: hashicorp/nomad
File: client/allocrunner/alloc_runner_test.go (Go)

Code before:
			return false, fmt.Errorf("got status %v; want %v", last.ClientStatus, structs.AllocClientStatusRunning)
		}
		return true, nil
	}, func(err error) {
		require.Fail(t, "err: %v", err)
	})

	update := ar.Alloc().Copy()
	migrate := true
	update.DesiredTransition.Migrate = &migrate
	update.DesiredStatus = structs.AllocDesiredStatusStop
	ar.Update(update)

	testutil.WaitForResult(func() (bool, error) {
		last := upd.Last()
		if last == nil {
			return false, fmt.Errorf("No updates")
		}

		if last.ClientStatus != structs.AllocClientStatusFailed {
			return false, fmt.Errorf("got client status %q; want %q", last.ClientStatus, structs.AllocClientStatusFailed)
		}

		return true, nil
	}, func(err error) {
		require.Fail(t, "err: %v", err)
	})
}

// Test that alloc runner kills tasks in task group when stopping and
// does not fail tasks when job is batch job type and not migrating
func TestAllocRunner_Batch_KillTG(t *testing.T) {
	ci.Parallel(t)

	alloc := mock.BatchAlloc()
	tr := alloc.AllocatedResources.Tasks[alloc.Job.TaskGroups[0].Tasks[0].Name]
	alloc.Job.TaskGroups[0].RestartPolicy.Attempts = 0
	alloc.Job.TaskGroups[0].Tasks[0].RestartPolicy.Attempts = 0

	task := alloc.Job.TaskGroups[0].Tasks[0]
	task.Driver = "mock_driver"
	task.Config["run_for"] = "10s"
	alloc.AllocatedResources.Tasks[task.Name] = tr

	conf, cleanup := testAllocRunnerConfig(t, alloc)
	defer cleanup()
	ar, err := NewAllocRunner(conf)
	must.NoError(t, err)

	defer destroy(ar)
	go ar.Run()

Review comment (suggestion):
Let's try to avoid adding more cases of `testify`:

```suggestion
		must.NoError(t, err)
```

Code after:

	// Wait for running
	testutil.WaitForResult(func() (bool, error) {
		last := upd.Last()
		if last == nil {
			return false, fmt.Errorf("No updates")
		}
		if last.ClientStatus != structs.AllocClientStatusRunning {
			return false, fmt.Errorf("got status %v; want %v", last.ClientStatus, structs.AllocClientStatusRunning)
		}
		return true, nil
	}, func(err error) {
		must.NoError(t, err)
	})

	// Wait for completed task
	testutil.WaitForResult(func() (bool, error) {
		last := upd.Last()
		if last == nil {
			return false, fmt.Errorf("No updates")
		}
		if last.ClientStatus != structs.AllocClientStatusRunning {
			return false, fmt.Errorf("got status %v; want %v", last.ClientStatus, structs.AllocClientStatusRunning)
		}

		// task should not have finished yet, task2 should be finished
		if !last.TaskStates[task.Name].FinishedAt.IsZero() {
			return false, fmt.Errorf("task should not be finished")
		}
		if last.TaskStates[task2.Name].FinishedAt.IsZero() {
			return false, fmt.Errorf("task should be finished")
		}
		return true, nil
	}, func(err error) {
		must.NoError(t, err)
	})

	update := ar.Alloc().Copy()
	migrate := true
	update.DesiredTransition.Migrate = &migrate
	update.DesiredStatus = structs.AllocDesiredStatusStop
	ar.Update(update)

	testutil.WaitForResult(func() (bool, error) {
		last := upd.Last()
		if last == nil {
			return false, fmt.Errorf("No updates")
		}

		if last.ClientStatus != structs.AllocClientStatusFailed {
			return false, fmt.Errorf("got client status %q; want %q", last.ClientStatus, structs.AllocClientStatusFailed)