From 1b2ff1013678bca2181d046a91857e78f3981f9c Mon Sep 17 00:00:00 2001
From: Roger Peppe <rogpeppe@gmail.com>
Date: Wed, 16 Oct 2019 20:50:17 +0100
Subject: [PATCH] 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>
---
 doc/go1.14.html         | 10 +++++++
 src/testing/sub_test.go | 65 +++++++++++++++++++++++++++++++++++++++++
 src/testing/testing.go  | 31 +++++++++++++++++++-
 3 files changed, 105 insertions(+), 1 deletion(-)

diff --git a/doc/go1.14.html b/doc/go1.14.html
index 58210b6529..37b14a50f0 100644
--- 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 abaedefde7..3f0f71f647 100644
--- 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 b9d4f2b5a5..59128e8a29 100644
--- 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()
-- 
2.50.0