Fix log callers when the cleanup uses t.Helper() (#19)

If every function in `t.Cleanup` uses `t.Helper()` and logs, then the
caller information skips to the caller of `t.Cleanup`. Track the caller
of `Cleanup` and use that instead of reporting faketb as caller.

TODO.md

@@ -2,6 +2,5 @@
 
 ## Caller information
 
- * Check how caller information is reported when `t.Helper` and `t.Cleanup` are mixed.
  * Provide ability to get structured information for logs which include caller, function, etc.
    * Clean up the API around gettling logs (testing log output format, list, string, etc).

fake_tb.go

@@ -8,6 +8,11 @@ import (
 	"testing"
 )
 
+const (
+	withSelf = 0
+	skipSelf = 1
+)
+
 var _ testing.TB = (*fakeTB)(nil)
 
 type fakeTB struct {
@@ -17,7 +22,7 @@ type fakeTB struct {
 
 	mu sync.Mutex // protects all of the below fields.
 
-	cleanups []func()
+	cleanups []cleanup
 	helpers  map[uintptr]struct{}
 	Logs     []logEntry
 
@@ -26,14 +31,24 @@ type fakeTB struct {
 	skipped   bool
 	panicked  bool
 
+	// only set during a cleanup
+	cleanupRoot  string
+	curCleanupPC []uintptr
+
 	// panic metadata
 	recovered      any
 	recoverCallers []uintptr
 }
 
 type logEntry struct {
-	callers []uintptr // callers[0] is the tb function that logged
-	entry   string
+	callers        []uintptr // callers[0] is the tb function that logged
+	cleanupCallers []uintptr // for logs within a cleanup function
+	entry          string
+}
+
+type cleanup struct {
+	fn      func()
+	callers []uintptr
 }
 
 // RunTest runs the given test using a fake [testing.TB] and returns
@@ -63,7 +78,7 @@ func (tb *fakeTB) checkPanic() {
 	if r := recover(); r != nil {
 		tb.panicked = true
 		tb.recovered = r
-		tb.recoverCallers = getCallers()
+		tb.recoverCallers = getCallers(skipSelf)
 	}
 }
 
@@ -73,6 +88,17 @@ func (tb *fakeTB) postTest() {
 }
 
 func (tb *fakeTB) runCleanups() {
+	// Set cleanupRoot so log callers can use cleanup's callers.
+	if self := getCaller(withSelf); self != 0 {
+		f := pcToFunction(self)
+		func() {
+			tb.mu.Lock()
+			defer tb.mu.Unlock()
+
+			tb.cleanupRoot = f
+		}()
+	}
+
 	// If defer runs with !finished, then a cleanup must have panicked
 	// (which could be a Skip/Fatal). Continue running remaining cleanups.
 	var finished bool
@@ -85,25 +111,33 @@ func (tb *fakeTB) runCleanups() {
 	// Run cleanups in last-first order, similar to defers.
 	// Don't iterate by index, as the slice can grow (cleanups can add cleanups).
 	for {
-		cleanupFn := func() func() {
+		c, ok := func() (cleanup, bool) {
 			tb.mu.Lock()
 			defer tb.mu.Unlock()
 
 			if len(tb.cleanups) == 0 {
-				return nil
+				return cleanup{}, false
 			}
 
 			last := len(tb.cleanups) - 1
-			fn := tb.cleanups[last]
+			c := tb.cleanups[last]
 			tb.cleanups = tb.cleanups[:last]
-			return fn
+			return c, true
 		}()
-
-		if cleanupFn == nil {
+		if !ok {
 			finished = true
 			break
 		}
-		cleanupFn()
+
+		// Set the caller of cleanup for logs with all `t.Helper()` frames.
+		func() {
+			tb.mu.Lock()
+			defer tb.mu.Unlock()
+
+			tb.curCleanupPC = c.callers
+		}()
+
+		c.fn()
 	}
 }
 
@@ -117,10 +151,14 @@ func (tb *fakeTB) done() bool {
 }
 
 func (tb *fakeTB) Cleanup(f func()) {
+	callerPCs := getCallers(skipSelf)
 	tb.mu.Lock()
 	defer tb.mu.Unlock()
 
-	tb.cleanups = append(tb.cleanups, f)
+	tb.cleanups = append(tb.cleanups, cleanup{
+		callers: callerPCs,
+		fn:      f,
+	})
 }
 
 func (tb *fakeTB) logfLocked(callers []uintptr, format string, args ...interface{}) {
@@ -138,24 +176,25 @@ func (tb *fakeTB) logLocked(callers []uintptr, args ...interface{}) {
 	formatted = strings.TrimSuffix(formatted, "\n")
 
 	tb.Logs = append(tb.Logs, logEntry{
-		callers: callers,
-		entry:   formatted,
+		callers:        callers,
+		cleanupCallers: tb.curCleanupPC,
+		entry:          formatted,
 	})
 }
 
 func (tb *fakeTB) Error(args ...interface{}) {
 	tb.mu.Lock()
 	defer tb.mu.Unlock()
 
-	tb.logLocked(getCallers(), args...)
+	tb.logLocked(getCallers(withSelf), args...)
 	tb.failLocked()
 }
 
 func (tb *fakeTB) Errorf(format string, args ...interface{}) {
 	tb.mu.Lock()
 	defer tb.mu.Unlock()
 
-	tb.logfLocked(getCallers(), format, args...)
+	tb.logfLocked(getCallers(withSelf), format, args...)
 	tb.failLocked()
 }
 
@@ -196,39 +235,37 @@ func (tb *fakeTB) Fatal(args ...interface{}) {
 	tb.mu.Lock()
 	defer tb.mu.Unlock()
 
-	tb.logLocked(getCallers(), args...)
+	tb.logLocked(getCallers(withSelf), args...)
 	tb.failNowLocked()
 }
 
 func (tb *fakeTB) Fatalf(format string, args ...interface{}) {
-	tb.logfLocked(getCallers(), format, args...)
+	tb.logfLocked(getCallers(withSelf), format, args...)
 	tb.FailNow()
 }
 
 func (tb *fakeTB) Helper() {
-	tb.mu.Lock()
-	defer tb.mu.Unlock()
-
-	const skip = 2 // runtime.Callers + Helper
-	var pc [1]uintptr
-	n := runtime.Callers(skip, pc[:])
-	if n == 0 {
+	callerPC := getCaller(skipSelf)
+	if callerPC == 0 {
 		// no callers, ignore.
 		// Note: real testing.TB would panic here, but we avoid panics in faket.
 		return
 	}
 
-	if _, ok := tb.helpers[pc[0]]; !ok {
-		tb.helpers[pc[0]] = struct{}{}
+	tb.mu.Lock()
+	defer tb.mu.Unlock()
+
+	if _, ok := tb.helpers[callerPC]; !ok {
+		tb.helpers[callerPC] = struct{}{}
 	}
 }
 
 func (tb *fakeTB) Log(args ...interface{}) {
-	tb.logLocked(getCallers(), args...)
+	tb.logLocked(getCallers(withSelf), args...)
 }
 
 func (tb *fakeTB) Logf(format string, args ...interface{}) {
-	tb.logfLocked(getCallers(), format, args...)
+	tb.logfLocked(getCallers(withSelf), format, args...)
 }
 
 func (tb *fakeTB) Name() string {
@@ -240,12 +277,12 @@ func (tb *fakeTB) Setenv(key, value string) {
 }
 
 func (tb *fakeTB) Skip(args ...interface{}) {
-	tb.logLocked(getCallers(), args...)
+	tb.logLocked(getCallers(withSelf), args...)
 	tb.skipNowLocked()
 }
 
 func (tb *fakeTB) Skipf(format string, args ...interface{}) {
-	tb.logfLocked(getCallers(), format, args...)
+	tb.logfLocked(getCallers(withSelf), format, args...)
 	tb.SkipNow()
 }
 
@@ -269,15 +306,32 @@ func (tb *fakeTB) TempDir() string {
 	return "tmp"
 }
 
-func getCallers() []uintptr {
+func getCallers(skip int) []uintptr {
+	skip += 2 // skip runtime.Callers and self.
 	depth := 32
 	for {
 		pc := make([]uintptr, depth)
-		// runtime.Callers returns itself, so skip that, and this function.
-		n := runtime.Callers(2, pc)
+		n := runtime.Callers(skip, pc)
 		if n < len(pc) {
 			return pc[:n]
 		}
 		depth *= 2
 	}
 }
+
+func getCaller(skip int) uintptr {
+	skip += 2 // skip runtime.Callers and this function
+	var pc [1]uintptr
+	n := runtime.Callers(skip, pc[:])
+	if n == 0 {
+		return 0
+	}
+
+	return pc[0]
+}
+
+func pcToFunction(pc uintptr) string {
+	frames := runtime.CallersFrames([]uintptr{pc})
+	f, _ := frames.Next()
+	return f.Function
+}

integration_test.go

@@ -217,3 +217,18 @@ func TestCmp_NestedCleanup(t *testing.T) {
 		})
 	}
 }
+
+func TestCmp_CleanupHelper(t *testing.T) {
+	cmptest.Compare(t, func(t testing.TB) {
+		helperWithCleanup(t)
+	})
+}
+
+func helperWithCleanup(t testing.TB) {
+	t.Helper()
+
+	t.Cleanup(func() {
+		t.Helper()
+		t.Log("cleanup func log")
+	})
+}

test_result.go

@@ -66,7 +66,7 @@ func (r TestResult) Logs() string {
 func (r TestResult) LogsWithCaller() []string {
 	logs := make([]string, 0, len(r.res.Logs))
 	for _, l := range r.res.Logs {
-		ci := getCallerInfo(l.callers, r.Helpers())
+		ci := getCallerInfo(l, r.res.cleanupRoot, r.Helpers())
 		line := fmt.Sprintf("%s:%d: %v", filepath.Base(ci.callerFile), ci.callerLine, l.entry)
 		logs = append(logs, line)
 	}
@@ -78,10 +78,8 @@ func (r TestResult) LogsWithCaller() []string {
 func (r TestResult) Helpers() []string {
 	funcs := make([]string, 0, len(r.res.helpers))
 	for pc := range r.res.helpers {
-		frames := runtime.CallersFrames([]uintptr{pc})
-		f, _ := frames.Next()
-		if f != (runtime.Frame{}) {
-			funcs = append(funcs, f.Function)
+		if fn := pcToFunction(pc); fn != "" {
+			funcs = append(funcs, fn)
 		}
 	}
 	sort.Strings(funcs)
@@ -95,7 +93,7 @@ type callerInfo struct {
 	callerLine int
 }
 
-func getCallerInfo(callers []uintptr, skipFuncs []string) callerInfo {
+func getCallerInfo(ent logEntry, cleanupRoot string, skipFuncs []string) callerInfo {
 	skipSet := make(map[string]struct{})
 	for _, f := range skipFuncs {
 		skipSet[f] = struct{}{}
@@ -105,7 +103,7 @@ func getCallerInfo(callers []uintptr, skipFuncs []string) callerInfo {
 	// but panic is not shown as a log caller.
 	skipSet["runtime.gopanic"] = struct{}{}
 
-	frames := runtime.CallersFrames(callers)
+	frames := runtime.CallersFrames(ent.callers)
 
 	f, _ := frames.Next()
 	if f == (runtime.Frame{}) {
@@ -124,6 +122,12 @@ func getCallerInfo(callers []uintptr, skipFuncs []string) callerInfo {
 			return ci
 		}
 
+		// If we hit the cleanup root, then use the callers of the t.Cleanup.
+		if f.Function == cleanupRoot {
+			frames = runtime.CallersFrames(ent.cleanupCallers)
+			continue
+		}
+
 		_, skip = skipSet[f.Function]
 	}
 

testdata/cmp_test_results.json

@@ -135,6 +135,11 @@
 {"Time":"2022-06-11T00:00:00.0Z","Action":"output","Package":"github.com/prashantv/faket","Test":"TestCmp_NestedCleanup/mix_nesting_and_skips","Output":"    --- SKIP: TestCmp_NestedCleanup/mix_nesting_and_skips (0.01s)\n"}
 {"Time":"2022-06-11T00:00:00.0Z","Action":"skip","Package":"github.com/prashantv/faket","Test":"TestCmp_NestedCleanup/mix_nesting_and_skips","Elapsed":0}
 {"Time":"2022-06-11T00:00:00.0Z","Action":"pass","Package":"github.com/prashantv/faket","Test":"TestCmp_NestedCleanup","Elapsed":0}
+{"Time":"2022-06-11T00:00:00.0Z","Action":"run","Package":"github.com/prashantv/faket","Test":"TestCmp_CleanupHelper"}
+{"Time":"2022-06-11T00:00:00.0Z","Action":"output","Package":"github.com/prashantv/faket","Test":"TestCmp_CleanupHelper","Output":"=== RUN   TestCmp_CleanupHelper\n"}
+{"Time":"2022-06-11T00:00:00.0Z","Action":"output","Package":"github.com/prashantv/faket","Test":"TestCmp_CleanupHelper","Output":"    integration_test.go:223: cleanup func log\n"}
+{"Time":"2022-06-11T00:00:00.0Z","Action":"output","Package":"github.com/prashantv/faket","Test":"TestCmp_CleanupHelper","Output":"--- PASS: TestCmp_CleanupHelper (0.01s)\n"}
+{"Time":"2022-06-11T00:00:00.0Z","Action":"pass","Package":"github.com/prashantv/faket","Test":"TestCmp_CleanupHelper","Elapsed":0}
 {"Time":"2022-06-11T00:00:00.0Z","Action":"output","Package":"github.com/prashantv/faket","Output":"FAIL\n"}
 {"Time":"2022-06-11T00:00:00.0Z","Action":"output","Package":"github.com/prashantv/faket","Output":"exit status 1\n"}
 {"Time":"2022-06-11T00:00:00.0Z","Action":"output","Package":"github.com/prashantv/faket","Output":"FAIL\tgithub.com/prashantv/faket\t0.01s\n"}