From ef77b645b927bcf4c049363c2e88a9f5c16446dd Mon Sep 17 00:00:00 2001
From: Dan Kegel <dkegel@fastly.com>
Date: Sun, 12 Dec 2021 12:29:27 -0800
Subject: [PATCH] testing: implement testing.Cleanup

Also reorder and regroup common's fields slightly to match upstream.

TODO: pull in more upstream tests once this package is goroutine-safe
---
 src/testing/sub_test.go | 80 +++++++++++++++++++++++++++++++++++++++++
 src/testing/testing.go  | 46 +++++++++++++++++++-----
 2 files changed, 117 insertions(+), 9 deletions(-)
 create mode 100644 src/testing/sub_test.go

diff --git a/src/testing/sub_test.go b/src/testing/sub_test.go
new file mode 100644
index 00000000..cd50b5b9
--- /dev/null
+++ b/src/testing/sub_test.go
@@ -0,0 +1,80 @@
+// Copyright 2016 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package testing
+
+import (
+	"reflect"
+)
+
+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 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)
+	}
+}
+
+func TestCleanupParallelSubtests(t *T) {
+	ranCleanup := 0
+	t.Run("test", func(t *T) {
+		t.Cleanup(func() { ranCleanup++ })
+		t.Run("x", func(t *T) {
+			t.Parallel()
+			if ranCleanup > 0 {
+				t.Error("outer cleanup ran before parallel subtest")
+			}
+		})
+	})
+	if ranCleanup != 1 {
+		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 5ff85e2b..b3ca8b30 100644
--- a/src/testing/testing.go
+++ b/src/testing/testing.go
@@ -43,16 +43,17 @@ func Init() {
 // common holds the elements common between T and B and
 // captures common methods such as Errorf.
 type common struct {
-	output bytes.Buffer
-	w      io.Writer // either &output, or at top level, os.Stdout
-	indent string
+	output   bytes.Buffer
+	w        io.Writer // either &output, or at top level, os.Stdout
+	indent   string
+	failed   bool     // Test or benchmark has failed.
+	skipped  bool     // Test of benchmark has been skipped.
+	cleanups []func() // optional functions to be called at the end of the test
+	finished bool     // Test function has completed.
 
-	failed   bool   // Test or benchmark has failed.
-	skipped  bool   // Test of benchmark has been skipped.
-	finished bool   // Test function has completed.
-	level    int    // Nesting depth of test or benchmark.
-	name     string // Name of test or benchmark.
-	parent   *common
+	parent *common
+	level  int    // Nesting depth of test or benchmark.
+	name   string // Name of test or benchmark.
 }
 
 // TB is the interface common to T and B.
@@ -208,7 +209,34 @@ func (c *common) Parallel() {
 	// Unimplemented.
 }
 
+// Cleanup registers a function to be called when the test (or subtest) and all its
+// subtests complete. Cleanup functions will be called in last added,
+// first called order.
+func (c *common) Cleanup(f func()) {
+	c.cleanups = append(c.cleanups, f)
+}
+
+// runCleanup is called at the end of the test.
+func (c *common) runCleanup() {
+	for {
+		var cleanup func()
+		if len(c.cleanups) > 0 {
+			last := len(c.cleanups) - 1
+			cleanup = c.cleanups[last]
+			c.cleanups = c.cleanups[:last]
+		}
+		if cleanup == nil {
+			return
+		}
+		cleanup()
+	}
+}
+
 func tRunner(t *T, fn func(t *T)) {
+	defer func() {
+		t.runCleanup()
+	}()
+
 	// Run the test.
 	if flagVerbose {
 		fmt.Fprintf(t.w, "=== RUN   %s\n", t.name)