Merge pull request #180 from dnephin/cleanup

Add calls to testing.TB.Cleanup() when available.

.golangci.yml

@@ -37,6 +37,8 @@ issues:
       text: .
     - text: 'return value of .*Close` is not checked'
       linters: [errcheck]
+    - text: 'SA1019'
+      linters: [staticcheck]
 
 linters:
   enable-all: true
@@ -47,3 +49,4 @@ linters:
     - gosec
     - scopelint
     - maligned
+    - gocritic

env/env.go

@@ -8,7 +8,7 @@ import (
 	"strings"
 
 	"gotest.tools/v3/assert"
-	"gotest.tools/v3/x/subtest"
+	"gotest.tools/v3/internal/cleanup"
 )
 
 type helperT interface {
@@ -18,30 +18,34 @@ type helperT interface {
 // Patch changes the value of an environment variable, and returns a
 // function which will reset the the value of that variable back to the
 // previous state.
+//
+// When used with Go 1.14+ the unpatch function will be called automatically
+// when the test ends, unless the TEST_NOCLEANUP env var is set to true.
 func Patch(t assert.TestingT, key, value string) func() {
 	if ht, ok := t.(helperT); ok {
 		ht.Helper()
 	}
-	oldValue, ok := os.LookupEnv(key)
+	oldValue, envVarExists := os.LookupEnv(key)
 	assert.NilError(t, os.Setenv(key, value))
-	cleanup := func() {
+	clean := func() {
 		if ht, ok := t.(helperT); ok {
 			ht.Helper()
 		}
-		if !ok {
+		if !envVarExists {
 			assert.NilError(t, os.Unsetenv(key))
 			return
 		}
 		assert.NilError(t, os.Setenv(key, oldValue))
 	}
-	if tc, ok := t.(subtest.TestContext); ok {
-		tc.AddCleanup(cleanup)
-	}
-	return cleanup
+	cleanup.Cleanup(t, clean)
+	return clean
 }
 
 // PatchAll sets the environment to env, and returns a function which will
 // reset the environment back to the previous state.
+//
+// When used with Go 1.14+ the unpatch function will be called automatically
+// when the test ends, unless the TEST_NOCLEANUP env var is set to true.
 func PatchAll(t assert.TestingT, env map[string]string) func() {
 	if ht, ok := t.(helperT); ok {
 		ht.Helper()
@@ -52,7 +56,7 @@ func PatchAll(t assert.TestingT, env map[string]string) func() {
 	for key, value := range env {
 		assert.NilError(t, os.Setenv(key, value), "setenv %s=%s", key, value)
 	}
-	cleanup := func() {
+	clean := func() {
 		if ht, ok := t.(helperT); ok {
 			ht.Helper()
 		}
@@ -61,10 +65,8 @@ func PatchAll(t assert.TestingT, env map[string]string) func() {
 			assert.NilError(t, os.Setenv(key, oldVal), "setenv %s=%s", key, oldVal)
 		}
 	}
-	if tc, ok := t.(subtest.TestContext); ok {
-		tc.AddCleanup(cleanup)
-	}
-	return cleanup
+	cleanup.Cleanup(t, clean)
+	return clean
 }
 
 // ToMap takes a list of strings in the format returned by os.Environ() and
@@ -94,21 +96,23 @@ func getParts(raw string) (string, string) {
 
 // ChangeWorkingDir to the directory, and return a function which restores the
 // previous working directory.
+//
+// When used with Go 1.14+ the previous working directory will be restored
+// automatically when the test ends, unless the TEST_NOCLEANUP env var is set to
+// true.
 func ChangeWorkingDir(t assert.TestingT, dir string) func() {
 	if ht, ok := t.(helperT); ok {
 		ht.Helper()
 	}
 	cwd, err := os.Getwd()
 	assert.NilError(t, err)
 	assert.NilError(t, os.Chdir(dir))
-	cleanup := func() {
+	clean := func() {
 		if ht, ok := t.(helperT); ok {
 			ht.Helper()
 		}
 		assert.NilError(t, os.Chdir(cwd))
 	}
-	if tc, ok := t.(subtest.TestContext); ok {
-		tc.AddCleanup(cleanup)
-	}
-	return cleanup
+	cleanup.Cleanup(t, clean)
+	return clean
 }

fs/file.go

@@ -11,7 +11,7 @@ import (
 	"strings"
 
 	"gotest.tools/v3/assert"
-	"gotest.tools/v3/x/subtest"
+	"gotest.tools/v3/internal/cleanup"
 )
 
 // Path objects return their filesystem path. Path may be implemented by a
@@ -36,27 +36,23 @@ type helperT interface {
 	Helper()
 }
 
-type cleanupT interface {
-	Cleanup(f func())
-}
-
 // NewFile creates a new file in a temporary directory using prefix as part of
 // the filename. The PathOps are applied to the before returning the File.
+//
+// When used with Go 1.14+ the file will be automatically removed when the test
+// ends, unless the TEST_NOCLEANUP env var is set to true.
 func NewFile(t assert.TestingT, prefix string, ops ...PathOp) *File {
 	if ht, ok := t.(helperT); ok {
 		ht.Helper()
 	}
 	tempfile, err := ioutil.TempFile("", cleanPrefix(prefix)+"-")
 	assert.NilError(t, err)
+
 	file := &File{path: tempfile.Name()}
+	cleanup.Cleanup(t, file.Remove)
+
 	assert.NilError(t, tempfile.Close())
 	assert.NilError(t, applyPathOps(file, ops))
-	if tc, ok := t.(subtest.TestContext); ok {
-		tc.AddCleanup(file.Remove)
-	}
-	if ct, ok := t.(cleanupT); ok {
-		ct.Cleanup(file.Remove)
-	}
 	return file
 }
 
@@ -86,20 +82,19 @@ type Dir struct {
 
 // NewDir returns a new temporary directory using prefix as part of the directory
 // name. The PathOps are applied before returning the Dir.
+//
+// When used with Go 1.14+ the directory will be automatically removed when the test
+// ends, unless the TEST_NOCLEANUP env var is set to true.
 func NewDir(t assert.TestingT, prefix string, ops ...PathOp) *Dir {
 	if ht, ok := t.(helperT); ok {
 		ht.Helper()
 	}
 	path, err := ioutil.TempDir("", cleanPrefix(prefix)+"-")
 	assert.NilError(t, err)
 	dir := &Dir{path: path}
+	cleanup.Cleanup(t, dir.Remove)
+
 	assert.NilError(t, applyPathOps(dir, ops))
-	if tc, ok := t.(subtest.TestContext); ok {
-		tc.AddCleanup(dir.Remove)
-	}
-	if ct, ok := t.(cleanupT); ok {
-		ct.Cleanup(dir.Remove)
-	}
 	return dir
 }
 

internal/cleanup/cleanup.go

@@ -0,0 +1,45 @@
+/*Package cleanup handles migration to and support for the Go 1.14+
+testing.TB.Cleanup() function.
+*/
+package cleanup
+
+import (
+	"os"
+	"strings"
+
+	"gotest.tools/v3/x/subtest"
+)
+
+type cleanupT interface {
+	Cleanup(f func())
+}
+
+type logT interface {
+	Log(...interface{})
+}
+
+type helperT interface {
+	Helper()
+}
+
+var noCleanup = strings.ToLower(os.Getenv("TEST_NOCLEANUP")) == "true"
+
+// Cleanup registers f as a cleanup function on t if any mechanisms are available.
+//
+// Skips registering f if TEST_NOCLEANUP is set to true.
+func Cleanup(t logT, f func()) {
+	if ht, ok := t.(helperT); ok {
+		ht.Helper()
+	}
+	if noCleanup {
+		t.Log("skipping cleanup because TEST_NOCLEANUP was enabled.")
+		return
+	}
+	if ct, ok := t.(cleanupT); ok {
+		ct.Cleanup(f)
+		return
+	}
+	if tc, ok := t.(subtest.TestContext); ok {
+		tc.AddCleanup(f)
+	}
+}

x/doc.go

@@ -1,4 +1,4 @@
-/*Package x is a namespace for other packages. Packages under x have looser
+/*Package x is a namespace for experimental packages. Packages under x have looser
 compatibility requirements. Packages in this namespace may contain backwards
 incompatible changes within the same major version.
 */

x/subtest/context.go

@@ -68,6 +68,9 @@ func Run(t *testing.T, name string, subtest func(t TestContext)) bool {
 type TestContext interface {
 	testing.TB
 	// AddCleanup function which will be run when before Run returns.
+	//
+	// Deprecated: Go 1.14+ now includes a testing.TB.Cleanup(func()) which
+	// should be used instead. AddCleanup will be removed in a future release.
 	AddCleanup(f func())
 	// Ctx returns a context for the test case. Multiple calls from the same subtest
 	// will return the same context. The context is cancelled when Run