container/stub: override goexit methods

FailNow, Fatal, Fatalf, SkipNow, Skip and Skipf must be called from the goroutine created by the test.

Signed-off-by: Ophestra <cat@gensokyo.uk>

1 files changed, 27 insertions, 6 deletions

diff --git a/container/stub/exit.go b/container/stub/exit.go
index 8d859461..0f44b5dc 100644
--- a/container/stub/exit.go
+++ b/container/stub/exit.go
@@ -1,15 +1,36 @@
 package stub
 
+import "testing"
+
 // PanicExit is a magic panic value treated as a simulated exit.
 const PanicExit = 0xdeadbeef
 
+const (
+	panicFailNow = 0xcafe0000 + iota
+	panicFatal
+	panicFatalf
+)
+
 // HandleExit must be deferred before calling with the stub.
-func HandleExit() {
-	r := recover()
-	if r == PanicExit {
-		return
-	}
-	if r != nil {
+func (s *Stub[K]) HandleExit() { handleExit(s.TB, true) }
+
+func handleExit(t testing.TB, root bool) {
+	switch r := recover(); r {
+	case PanicExit:
+		break
+
+	case panicFailNow:
+		if root {
+			t.FailNow()
+		} else {
+			t.Fail()
+		}
+		break
+
+	case panicFatal, panicFatalf, nil:
+		break
+
+	default:
 		panic(r)
 	}
 }