testing: make Cleanup work for benchmarks too.

Fixes #37073.

Change-Id: I6fb24a3f9d7b7adf3213ac6a8bcbf5fb43975b7e
Reviewed-on: https://go-review.googlesource.com/c/go/+/218117
Reviewed-by: Bryan C. Mills <bcmills@google.com>
Reviewed-by: Ian Lance Taylor <iant@golang.org>
Run-TryBot: Bryan C. Mills <bcmills@google.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>

diff --git a/src/testing/benchmark.go b/src/testing/benchmark.go
index 93f461b07a22d32e3ea6ee15ad077885de0b2679..88ba0f024203fa70d9fb5f97f74b179b36bd5d6d 100644 (file)
--- a/src/testing/benchmark.go
+++ b/src/testing/benchmark.go
@@ -179,6 +179,7 @@ func (b *B) ReportAllocs() {
 func (b *B) runN(n int) {
        benchmarkLock.Lock()
        defer benchmarkLock.Unlock()
+       defer b.runCleanup(normalPanic)
        // Try to get a comparable environment for each run
        // by clearing garbage from previous runs.
        runtime.GC()

diff --git a/src/testing/sub_test.go b/src/testing/sub_test.go
index 3dc30ee72e2077ecd79f078ccf1395978a0d1a93..95f8220f815fe08834e566224f67a5caee93df20 100644 (file)
--- a/src/testing/sub_test.go
+++ b/src/testing/sub_test.go
@@ -613,6 +613,46 @@ func TestBRun(t *T) {
                                t.Errorf("MemBytes was %v; want %v", got, 2*bufSize)
                        }
                },
+       }, {
+               desc: "cleanup is called",
+               f: func(b *B) {
+                       var calls, cleanups, innerCalls, innerCleanups int
+                       b.Run("", func(b *B) {
+                               calls++
+                               b.Cleanup(func() {
+                                       cleanups++
+                               })
+                               b.Run("", func(b *B) {
+                                       b.Cleanup(func() {
+                                               innerCleanups++
+                                       })
+                                       innerCalls++
+                               })
+                               work(b)
+                       })
+                       if calls == 0 || calls != cleanups {
+                               t.Errorf("mismatched cleanups; got %d want %d", cleanups, calls)
+                       }
+                       if innerCalls == 0 || innerCalls != innerCleanups {
+                               t.Errorf("mismatched cleanups; got %d want %d", cleanups, calls)
+                       }
+               },
+       }, {
+               desc:   "cleanup is called on failure",
+               failed: true,
+               f: func(b *B) {
+                       var calls, cleanups int
+                       b.Run("", func(b *B) {
+                               calls++
+                               b.Cleanup(func() {
+                                       cleanups++
+                               })
+                               b.Fatalf("failure")
+                       })
+                       if calls == 0 || calls != cleanups {
+                               t.Errorf("mismatched cleanups; got %d want %d", cleanups, calls)
+                       }
+               },
        }}
        for _, tc := range testCases {
                var ok bool