testing: implement Cleanup(), with tests

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)
+	}
+}

src/testing/testing.go

@@ -12,6 +12,7 @@ import (
 	"bytes"
 	"flag"
 	"fmt"
+	"io"
 	"os"
 	"strings"
 )
@@ -42,14 +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
-	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
+	level  int    // Nesting depth of test or benchmark.
+	name   string // Name of test or benchmark.
 }
 
 // TB is the interface common to T and B.
@@ -205,35 +209,70 @@ func (c *common) Parallel() {
 	// Unimplemented.
 }
 
-// Run runs a subtest of f t called name. It waits until the subtest is finished
-// and returns whether the subtest succeeded.
-func (t *T) Run(name string, f func(t *T)) bool {
-	// Create a subtest.
-	sub := T{
-		common: common{
-			name:   t.name + "/" + rewrite(name),
-			indent: t.indent + "    ",
-		},
+// 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.output, "=== RUN   %s\n", sub.name)
-
+		fmt.Fprintf(t.w, "=== RUN   %s\n", t.name)
 	}
-	f(&sub)
+
+	fn(t)
 
 	// Process the result (pass or fail).
-	if sub.failed {
-		t.failed = true
-		fmt.Fprintf(&t.output, sub.indent+"--- FAIL: %s\n", sub.name)
-		t.output.Write(sub.output.Bytes())
+	if t.failed {
+		if t.parent != nil {
+			t.parent.failed = true
+		}
+		fmt.Fprintf(t.w, t.indent+"--- FAIL: %s\n", t.name)
+		t.w.Write(t.output.Bytes())
 	} else {
 		if flagVerbose {
-			fmt.Fprintf(&t.output, sub.indent+"--- PASS: %s\n", sub.name)
-			t.output.Write(sub.output.Bytes())
+			fmt.Fprintf(t.w, t.indent+"--- PASS: %s\n", t.name)
+			t.w.Write(t.output.Bytes())
 		}
 	}
+}
+
+// Run runs f as a subtest of t called name. It waits until the subtest is finished
+// and returns whether the subtest succeeded.
+func (t *T) Run(name string, f func(t *T)) bool {
+	// Create a subtest.
+	sub := T{
+		common: common{
+			name:   t.name + "/" + rewrite(name),
+			indent: t.indent + "    ",
+			w:      &t.output,
+			parent: &t.common,
+		},
+	}
+
+	tRunner(&sub, f)
 	return !sub.failed
 }
 
@@ -313,23 +352,11 @@ func (m *M) Run() int {
 		t := &T{
 			common: common{
 				name: test.Name,
+				w:    os.Stdout,
 			},
 		}
 
-		if flagVerbose {
-			fmt.Printf("=== RUN   %s\n", test.Name)
-		}
-		test.F(t)
-
-		if t.failed {
-			fmt.Printf("--- FAIL: %s\n", test.Name)
-			os.Stdout.Write(t.output.Bytes())
-		} else {
-			if flagVerbose {
-				fmt.Printf("--- PASS: %s\n", test.Name)
-				os.Stdout.Write(t.output.Bytes())
-			}
-		}
+		tRunner(t, test.F)
 
 		if t.failed {
 			failures++

testdata/testing.go

@@ -3,6 +3,8 @@ package main
 // TODO: also test the verbose version.
 
 import (
+	"errors"
+	"io"
 	"testing"
 )
 
@@ -33,14 +35,26 @@ var benchmarks = []testing.InternalBenchmark{}
 
 var examples = []testing.InternalExample{}
 
+var errMain = errors.New("testing: unexpected use of func Main")
+
+// matchStringOnly is part of upstream, and is used below to provide a dummy deps to pass to MainStart
+// so it can be run with go (tested with go 1.16) to provide a baseline for the regression test.
+// See c56cc9b3b57276.  Unfortunately, testdeps is internal, so we can't just use &testdeps.TestDeps{}.
+type matchStringOnly func(pat, str string) (bool, error)
+
+func (f matchStringOnly) MatchString(pat, str string) (bool, error)   { return f(pat, str) }
+func (f matchStringOnly) StartCPUProfile(w io.Writer) error           { return errMain }
+func (f matchStringOnly) StopCPUProfile()                             {}
+func (f matchStringOnly) WriteProfileTo(string, io.Writer, int) error { return errMain }
+func (f matchStringOnly) ImportPath() string                          { return "" }
+func (f matchStringOnly) StartTestLog(io.Writer)                      {}
+func (f matchStringOnly) StopTestLog() error                          { return errMain }
+func (f matchStringOnly) SetPanicOnExit0(bool)                        {}
+
 func main() {
-	m := testing.MainStart(testdeps{}, tests, benchmarks, examples)
+	m := testing.MainStart(matchStringOnly(nil), tests, benchmarks, examples)
 	exitcode := m.Run()
 	if exitcode != 0 {
 		println("exitcode:", exitcode)
 	}
 }
-
-type testdeps struct{}
-
-func (testdeps) MatchString(pat, str string) (bool, error) { return true, nil }