testing: add TB.TempDir

Fixes #35998

Change-Id: I87c6bf4e34e832be68862ca16ecfa6ea12048d31
Reviewed-on: https://go-review.googlesource.com/c/go/+/226877
Run-TryBot: Brad Fitzpatrick <bradfitz@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Ian Lance Taylor <iant@golang.org>

2 files changed, 78 insertions, 0 deletions

diff --git a/src/testing/testing.go b/src/testing/testing.go
index 85a92c9384..d546f56478 100644
--- a/src/testing/testing.go
+++ b/src/testing/testing.go
@@ -239,6 +239,7 @@ import (
 	"fmt"
 	"internal/race"
 	"io"
+	"io/ioutil"
 	"os"
 	"runtime"
 	"runtime/debug"
@@ -362,6 +363,10 @@ type common struct {
 	barrier  chan bool // To signal parallel subtests they may start.
 	signal   chan bool // To signal a test is done.
 	sub      []*T      // Queue of subtests to be run in parallel.
+
+	tempDirOnce sync.Once
+	tempDir     string
+	tempDirErr  error
 }
 
 // Short reports whether the -test.short flag is set.
@@ -561,6 +566,7 @@ type TB interface {
 	SkipNow()
 	Skipf(format string, args ...interface{})
 	Skipped() bool
+	TempDir() string
 
 	// A private method to prevent users implementing the
 	// interface and so future additions to it will not
@@ -791,6 +797,30 @@ func (c *common) Cleanup(f func()) {
 	}
 }
 
+// TempDir returns a temporary directory for the test to use.
+// It is lazily created on first access, and calls t.Fatal if the directory
+// creation fails.
+// Subsequent calls to t.TempDir return the same directory.
+// The directory is automatically removed by Cleanup when the test and
+// all its subtests complete.
+func (c *common) TempDir() string {
+	c.tempDirOnce.Do(func() {
+		c.Helper()
+		c.tempDir, c.tempDirErr = ioutil.TempDir("", c.Name())
+		if c.tempDirErr == nil {
+			c.Cleanup(func() {
+				if err := os.RemoveAll(c.tempDir); err != nil {
+					c.Errorf("TempDir RemoveAll cleanup: %v", err)
+				}
+			})
+		}
+	})
+	if c.tempDirErr != nil {
+		c.Fatalf("TempDir: %v", c.tempDirErr)
+	}
+	return c.tempDir
+}
+
 // panicHanding is an argument to runCleanup.
 type panicHandling int
 
diff --git a/src/testing/testing_test.go b/src/testing/testing_test.go
index 45e44683b4..afb35a96d4 100644
--- a/src/testing/testing_test.go
+++ b/src/testing/testing_test.go
@@ -5,6 +5,7 @@
 package testing_test
 
 import (
+	"io/ioutil"
 	"os"
 	"testing"
 )
@@ -16,3 +17,50 @@ import (
 func TestMain(m *testing.M) {
 	os.Exit(m.Run())
 }
+
+func TestTempDir(t *testing.T) {
+	dirCh := make(chan string, 1)
+	t.Cleanup(func() {
+		// Verify directory has been removed.
+		select {
+		case dir := <-dirCh:
+			fi, err := os.Stat(dir)
+			if os.IsNotExist(err) {
+				// All good
+				return
+			}
+			if err != nil {
+				t.Fatal(err)
+			}
+			t.Errorf("directory %q stil exists: %v, isDir=%v", dir, fi, fi.IsDir())
+		default:
+			if !t.Failed() {
+				t.Fatal("never received dir channel")
+			}
+		}
+	})
+
+	dir := t.TempDir()
+	if dir == "" {
+		t.Fatal("expected dir")
+	}
+	dir2 := t.TempDir()
+	if dir != dir2 {
+		t.Fatal("directory changed between calls")
+	}
+	dirCh <- dir
+	fi, err := os.Stat(dir)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !fi.IsDir() {
+		t.Errorf("dir %q is not a dir", dir)
+	}
+	fis, err := ioutil.ReadDir(dir)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(fis) > 0 {
+		t.Errorf("unexpected %d files in TempDir: %v", len(fis), fis)
+	}
+}