testing: implement testing.Cleanup

TODO: pull in more upstream tests once this package is goroutine-safe

Author: dkegel-fastly

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

@@ -47,11 +47,12 @@ type common struct {
 	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.
-	finished bool   // Test function has completed.
-	name     string // Name of test or benchmark.
+	failed   bool     // Test or benchmark has failed.
+	skipped  bool     // Test of benchmark has been skipped.
+	finished bool     // Test function has completed.
+	cleanups []func() // optional functions to be called at the end of the test
 	parent   *common
+	name     string // Name of test or benchmark.
 }
 
 // TB is the interface common to T and B.
@@ -207,7 +208,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)