From e059adef98924174a43ce5d6b0a3f0665189e856 Mon Sep 17 00:00:00 2001
From: Daniel Bennett <dbennett@hashicorp.com>
Date: Mon, 29 Jan 2024 17:54:54 -0600
Subject: [PATCH] e2e: PreCleanup and other jobs3 helpers (#19844)

---
 e2e/v3/jobs3/jobs3.go | 84 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 84 insertions(+)

diff --git a/e2e/v3/jobs3/jobs3.go b/e2e/v3/jobs3/jobs3.go
index 52514975b..b2c872a0c 100644
--- a/e2e/v3/jobs3/jobs3.go
+++ b/e2e/v3/jobs3/jobs3.go
@@ -40,6 +40,8 @@ type Submission struct {
 
 	// jobspec mutator funcs
 	mutators []func(string) string
+	// preCleanup funcs to run before deregistering the job
+	preCleanup []func(*Submission)
 
 	vars         *set.Set[string] // key=value
 	waitComplete *set.Set[string] // groups to wait until complete
@@ -54,6 +56,45 @@ func (sub *Submission) queryOptions() *nomadapi.QueryOptions {
 	}
 }
 
+func (sub *Submission) Evals() []*nomadapi.Evaluation {
+	sub.t.Helper()
+	evals, _, err := sub.nomadClient.Jobs().
+		Evaluations(sub.JobID(), sub.queryOptions())
+	must.NoError(sub.t, err)
+	return evals
+}
+
+func (sub *Submission) Allocs() []*nomadapi.AllocationListStub {
+	sub.t.Helper()
+	allocs, _, err := sub.nomadClient.Jobs().
+		Allocations(sub.jobID, true, sub.queryOptions())
+	must.NoError(sub.t, err, must.Sprint("could not get allocs"))
+	return allocs
+}
+
+type TaskEvents struct {
+	Group  string
+	Task   string
+	Events []*nomadapi.TaskEvent
+}
+
+// AllocEvents returns a map of TaskEvents with alloc ID keys
+func (sub *Submission) AllocEvents() map[string]TaskEvents {
+	sub.t.Helper()
+	allocs := sub.Allocs()
+	events := make(map[string]TaskEvents)
+	for _, alloc := range allocs {
+		for task, state := range alloc.TaskStates {
+			events[alloc.ID] = TaskEvents{
+				Group:  alloc.TaskGroup,
+				Task:   task,
+				Events: state.Events,
+			}
+		}
+	}
+	return events
+}
+
 type Logs struct {
 	Stdout string
 	Stderr string
@@ -273,6 +314,14 @@ func (sub *Submission) run() {
 		sub.t.Cleanup(sub.cleanup)
 	}
 
+	// pre-cleanup callbacks run before main cleanup (reverse order of their
+	// addition with t.Cleanup())
+	for _, f := range sub.preCleanup {
+		sub.t.Cleanup(func() {
+			f(sub)
+		})
+	}
+
 	evalID := regResp.EvalID
 
 	queryOpts := &nomadapi.QueryOptions{
@@ -470,6 +519,7 @@ func initialize(t *testing.T, filename string) *Submission {
 		timeout:      20 * time.Second,
 		vars:         set.New[string](0),
 		waitComplete: set.New[string](0),
+		preCleanup:   []func(*Submission){defaultPreCleanup},
 	}
 }
 
@@ -531,6 +581,40 @@ func WaitComplete(group string) Option {
 	}
 }
 
+// PreCleanup runs a function after run has completed, before cleanup.
+func PreCleanup(cb func(*Submission)) Option {
+	return func(sub *Submission) {
+		sub.preCleanup = append(sub.preCleanup, cb)
+	}
+}
+
+// defaultPreCleanup looks for blocked evals, alloc errors, and task events
+// only when the test has failed.
+func defaultPreCleanup(job *Submission) {
+	if !job.t.Failed() {
+		return
+	}
+
+	for _, eval := range job.Evals() {
+		for group, block := range eval.FailedTGAllocs {
+			job.t.Logf("eval for tg '%s' failed; constraints: %+v",
+				group, block.ConstraintFiltered)
+		}
+	}
+
+	for _, alloc := range job.Allocs() {
+		job.t.Logf("tg '%s' alloc status '%s': %s",
+			alloc.TaskGroup, alloc.ClientStatus, alloc.ClientDescription)
+	}
+
+	for _, ae := range job.AllocEvents() {
+		for _, event := range ae.Events {
+			job.t.Logf("tg '%s' task '%s' event: %s",
+				ae.Group, ae.Task, event.DisplayMessage)
+		}
+	}
+}
+
 // SkipEvalComplete will skip waiting for the evaluation(s) to be complete.
 //
 // Implies SkipDeploymentHealthy.