testing: implement Cleanup method

Fixes #32111

Change-Id: I7078947889d1e126d9679fb28f27b3fa6ce133ef
Reviewed-on: https://go-review.googlesource.com/c/go/+/201359
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>

diff --git a/doc/go1.14.html b/doc/go1.14.html
index 58210b65294e6ac760d0cfe20b98aa8f557c8321..37b14a50f0e59be116345630ac0a3e6d46e3ab10 100644 (file)
--- a/doc/go1.14.html
+++ b/doc/go1.14.html
@@ -248,6 +248,16 @@ TODO
 
 </dl><!-- runtime -->
 
+<dl id="testing"><dt><a href="/pkg/testing/">testing</a></dt>
+  <dd>
+    <p><!-- CL 201359 -->
+       The testing package now supports cleanup functions, called after
+       a test or benchmark has finished, by calling
+       <a href="/pkg/testing#T.Cleanup"><code>T.Cleanup</code></a> or
+       <a href="/pkg/testing#B.Cleanup"><code>B.Cleanup</code></a> respectively.
+    </p>
+</dl><!-- testing -->
+
 <h3 id="minor_library_changes">Minor changes to the library</h3>
 
 <p>

diff --git a/src/testing/sub_test.go b/src/testing/sub_test.go
index abaedefde7d352de48a97212b5d2b6c04f9c3e8b..3f0f71f647d8c11229d8e1dd92b8fc4b3874b8bc 100644 (file)
--- a/src/testing/sub_test.go
+++ b/src/testing/sub_test.go
@@ -7,6 +7,7 @@ package testing
 import (
        "bytes"
        "fmt"
+       "reflect"
        "regexp"
        "runtime"
        "strings"
@@ -790,3 +791,67 @@ func TestBenchmark(t *T) {
                t.Errorf("want >5ms; got %v", time.Duration(res.NsPerOp()))
        }
 }
+
+func TestCleanup(t *T) {
+       var cleanups []int
+       t.Run("test", func(t *T) {
+               t.Cleanup(func() { cleanups = append(cleanups, 1) })
+               t.Cleanup(func() { cleanups = append(cleanups, 2) })
+       })
+       if got, want := cleanups, []int{2, 1}; !reflect.DeepEqual(got, want) {
+               t.Errorf("unexpected cleanup record; got %v want %v", got, want)
+       }
+}
+
+func TestConcurrentCleanup(t *T) {
+       cleanups := 0
+       t.Run("test", func(t *T) {
+               done := make(chan struct{})
+               for i := 0; i < 2; i++ {
+                       i := i
+                       go func() {
+                               t.Cleanup(func() {
+                                       cleanups |= 1 << i
+                               })
+                               done <- struct{}{}
+                       }()
+               }
+               <-done
+               <-done
+       })
+       if cleanups != 1|2 {
+               t.Errorf("unexpected cleanup; got %d want 3", cleanups)
+       }
+}
+
+func TestCleanupCalledEvenAfterGoexit(t *T) {
+       cleanups := 0
+       t.Run("test", func(t *T) {
+               t.Cleanup(func() {
+                       cleanups++
+               })
+               t.Cleanup(func() {
+                       runtime.Goexit()
+               })
+       })
+       if cleanups != 1 {
+               t.Errorf("unexpected cleanup count; got %d want 1", cleanups)
+       }
+}
+
+func TestRunCleanup(t *T) {
+       outerCleanup := 0
+       innerCleanup := 0
+       t.Run("test", func(t *T) {
+               t.Cleanup(func() { outerCleanup++ })
+               t.Run("x", func(t *T) {
+                       t.Cleanup(func() { innerCleanup++ })
+               })
+       })
+       if innerCleanup != 1 {
+               t.Errorf("unexpected inner cleanup count; got %d want 1", innerCleanup)
+       }
+       if outerCleanup != 1 {
+               t.Errorf("unexpected outer cleanup count; got %d want 0", outerCleanup)
+       }
+}

diff --git a/src/testing/testing.go b/src/testing/testing.go
index b9d4f2b5a564846b18b939cd74351aa0be753784..59128e8a29d0855169d57c45ae1c5762ad5486d7 100644 (file)
--- a/src/testing/testing.go
+++ b/src/testing/testing.go
@@ -344,6 +344,7 @@ type common struct {
        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
 
        chatty     bool   // A copy of the chatty flag.
        finished   bool   // Test function has completed.
@@ -543,6 +544,7 @@ func fmtDuration(d time.Duration) string {
 
 // TB is the interface common to T and B.
 type TB interface {
+       Cleanup(func())
        Error(args ...interface{})
        Errorf(format string, args ...interface{})
        Fail()
@@ -550,6 +552,7 @@ type TB interface {
        Failed() bool
        Fatal(args ...interface{})
        Fatalf(format string, args ...interface{})
+       Helper()
        Log(args ...interface{})
        Logf(format string, args ...interface{})
        Name() string
@@ -557,7 +560,6 @@ type TB interface {
        SkipNow()
        Skipf(format string, args ...interface{})
        Skipped() bool
-       Helper()
 
        // A private method to prevent users implementing the
        // interface and so future additions to it will not
@@ -774,6 +776,32 @@ func (c *common) Helper() {
        c.helpers[callerName(1)] = struct{}{}
 }
 
+// Cleanup registers a function to be called when the test finishes.
+// 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
+       c.cleanup = func() {
+               if oldCleanup != nil {
+                       defer oldCleanup()
+               }
+               f()
+       }
+}
+
+// runCleanup is called at the end of the test.
+func (c *common) runCleanup() {
+       c.mu.Lock()
+       cleanup := c.cleanup
+       c.cleanup = nil
+       c.mu.Unlock()
+       if cleanup != nil {
+               cleanup()
+       }
+}
+
 // callerName gives the function name (qualified with a package path)
 // for the caller after skip frames (where 0 means the current function).
 func callerName(skip int) string {
@@ -919,6 +947,7 @@ func tRunner(t *T, fn func(t *T)) {
                }
                t.signal <- signal
        }()
+       defer t.runCleanup()
 
        t.start = time.Now()
        t.raceErrors = -race.Errors()