From bfea70d5e402db048661b48ec7ff77cff7728be9 Mon Sep 17 00:00:00 2001
From: Seebs <seebs@molecula.com>
Date: Thu, 21 Jul 2022 17:52:58 -0500
Subject: [PATCH] improve post-hook behavior

We want to be able to register hooks which do cleanup, which may be
registered after the auditor cleanup check, which means that we
want LIFO order for post-hook cleanups.

We also want the test hook cleanup to be deferred, rather than
merely run after the tests are executed. Also, we have to extract
the result from running the test, then execute deferred things,
*then* call os.Exit, because os.Exit bypasses defers.
---
 testhook/hook.go | 35 +++++++++++++++++++++++++----------
 1 file changed, 25 insertions(+), 10 deletions(-)

diff --git a/testhook/hook.go b/testhook/hook.go
index 6ca2aacb8..dd561c896 100644
--- a/testhook/hook.go
+++ b/testhook/hook.go
@@ -24,6 +24,10 @@ func RegisterPostTestHook(fn Callback) {
 	mu.Lock()
 	defer mu.Unlock()
 	postHooks = append(postHooks, fn)
+	// ... but put it at the beginning of the list, so they're LIFO,
+	// so test hook pairs nest cleanly.
+	copy(postHooks[1:], postHooks)
+	postHooks[0] = fn
 }
 
 // RegisterPreTestHook registers a function to be called after tests
@@ -54,16 +58,27 @@ func RunTestsWithHooks(m *testing.M) {
 		fmt.Fprint(os.Stderr, "pre-hooks failed, aborting.\n")
 		os.Exit(ret)
 	}
-	ret = m.Run()
-	mu.Lock()
-	defer mu.Unlock()
-	for _, fn := range postHooks {
-		err := fn()
-		if err != nil {
-			fmt.Fprintf(os.Stderr, "post-hook failure: %v\n", err)
-			ret = 1
-		}
-	}
+
+	// This inner function lets us get a deferred run of our post-test
+	// hooks which necessarily succeeds, but can trap an error even
+	// from those and set the return value. You can't do this with a
+	// direct os.Exit() wrapper because defers don't run after os.Exit.
+	// Don't ask how many tries it took me to figure that out, I'll
+	// just cry.
+	func() {
+		defer func() {
+			mu.Lock()
+			defer mu.Unlock()
+			for _, fn := range postHooks {
+				err := fn()
+				if err != nil {
+					fmt.Fprintf(os.Stderr, "post-hook failure: %v\n", err)
+					ret = 1
+				}
+			}
+		}()
+		ret = m.Run()
+	}()
 	os.Exit(ret)
 }