testing: run a Cleanup registered by a Cleanup
Fixes #41085
Change-Id: Ieafc60cbc8e09f1935d38b1767b084d78dae5cb4
Reviewed-on: https://go-review.googlesource.com/c/go/+/251457
Run-TryBot: Ian Lance Taylor <iant@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Bryan C. Mills <bcmills@google.com>
diff --git a/src/testing/sub_test.go b/src/testing/sub_test.go
index 8eb0084..51fc0cc 100644
--- a/src/testing/sub_test.go
+++ b/src/testing/sub_test.go
@@ -928,3 +928,30 @@
t.Errorf("unexpected cleanup count; got %d want 1", ranCleanup)
}
}
+
+func TestNestedCleanup(t *T) {
+ ranCleanup := 0
+ t.Run("test", func(t *T) {
+ t.Cleanup(func() {
+ if ranCleanup != 2 {
+ t.Errorf("unexpected cleanup count in first cleanup: got %d want 2", ranCleanup)
+ }
+ ranCleanup++
+ })
+ t.Cleanup(func() {
+ if ranCleanup != 0 {
+ t.Errorf("unexpected cleanup count in second cleanup: got %d want 0", ranCleanup)
+ }
+ ranCleanup++
+ t.Cleanup(func() {
+ if ranCleanup != 1 {
+ t.Errorf("unexpected cleanup count in nested cleanup: got %d want 1", ranCleanup)
+ }
+ ranCleanup++
+ })
+ })
+ })
+ if ranCleanup != 3 {
+ t.Errorf("unexpected cleanup count: got %d want 3", ranCleanup)
+ }
+}
diff --git a/src/testing/testing.go b/src/testing/testing.go
index 0174396..f4f0060 100644
--- a/src/testing/testing.go
+++ b/src/testing/testing.go
@@ -403,7 +403,7 @@
skipped bool // Test of benchmark has been skipped.
done bool // Test is finished and all subtests have completed.
helpers map[string]struct{} // functions to be skipped when writing file/line info
- cleanup func() // optional function to be called at the end of the test
+ cleanups []func() // optional functions to be called at the end of the test
cleanupName string // Name of the cleanup function.
cleanupPc []uintptr // The stack trace at the point where Cleanup was called.
@@ -855,28 +855,31 @@
// subtests complete. Cleanup functions will be called in last added,
// first called order.
func (c *common) Cleanup(f func()) {
- c.mu.Lock()
- defer c.mu.Unlock()
- oldCleanup := c.cleanup
- oldCleanupPc := c.cleanupPc
- c.cleanup = func() {
- if oldCleanup != nil {
- defer func() {
- c.mu.Lock()
- c.cleanupPc = oldCleanupPc
- c.mu.Unlock()
- oldCleanup()
- }()
- }
- c.mu.Lock()
- c.cleanupName = callerName(0)
- c.mu.Unlock()
- f()
- }
var pc [maxStackLen]uintptr
// Skip two extra frames to account for this function and runtime.Callers itself.
n := runtime.Callers(2, pc[:])
- c.cleanupPc = pc[:n]
+ cleanupPc := pc[:n]
+
+ fn := func() {
+ defer func() {
+ c.mu.Lock()
+ defer c.mu.Unlock()
+ c.cleanupName = ""
+ c.cleanupPc = nil
+ }()
+
+ name := callerName(0)
+ c.mu.Lock()
+ c.cleanupName = name
+ c.cleanupPc = cleanupPc
+ c.mu.Unlock()
+
+ f()
+ }
+
+ c.mu.Lock()
+ defer c.mu.Unlock()
+ c.cleanups = append(c.cleanups, fn)
}
var tempDirReplacer struct {
@@ -934,22 +937,37 @@
// If catchPanic is true, this will catch panics, and return the recovered
// value if any.
func (c *common) runCleanup(ph panicHandling) (panicVal interface{}) {
- c.mu.Lock()
- cleanup := c.cleanup
- c.cleanup = nil
- c.mu.Unlock()
- if cleanup == nil {
- return nil
- }
-
if ph == recoverAndReturnPanic {
defer func() {
panicVal = recover()
}()
}
- cleanup()
- return nil
+ // Make sure that if a cleanup function panics,
+ // we still run the remaining cleanup functions.
+ defer func() {
+ c.mu.Lock()
+ recur := len(c.cleanups) > 0
+ c.mu.Unlock()
+ if recur {
+ c.runCleanup(normalPanic)
+ }
+ }()
+
+ for {
+ var cleanup func()
+ c.mu.Lock()
+ if len(c.cleanups) > 0 {
+ last := len(c.cleanups) - 1
+ cleanup = c.cleanups[last]
+ c.cleanups = c.cleanups[:last]
+ }
+ c.mu.Unlock()
+ if cleanup == nil {
+ return nil
+ }
+ cleanup()
+ }
}
// callerName gives the function name (qualified with a package path)