testutil: Add t.Cleanup to TempDir

TempDir registers a Cleanup so that the directory is always removed. To disable to cleanup, set the TEST_NOCLEANUP
env var.
This commit is contained in:
Daniel Nephin 2020-08-14 13:19:06 -04:00
parent 6b0ac22c1b
commit 2725513eea
2 changed files with 12 additions and 9 deletions

View File

@ -10,7 +10,6 @@ import (
"io" "io"
"math/rand" "math/rand"
"net/http/httptest" "net/http/httptest"
"os"
"path/filepath" "path/filepath"
"strconv" "strconv"
"testing" "testing"
@ -38,9 +37,6 @@ func init() {
rand.Seed(time.Now().UnixNano()) // seed random number generator rand.Seed(time.Now().UnixNano()) // seed random number generator
} }
// TempDir defines the base dir for temporary directories.
var TempDir = os.TempDir()
// TestAgent encapsulates an Agent with a default configuration and // TestAgent encapsulates an Agent with a default configuration and
// startup procedure suitable for testing. It panics if there are errors // startup procedure suitable for testing. It panics if there are errors
// during creation or startup instead of returning errors. It manages a // during creation or startup instead of returning errors. It manages a

View File

@ -29,21 +29,28 @@ func init() {
} }
} }
var noCleanup = strings.ToLower(os.Getenv("TEST_NOCLEANUP")) == "true"
// TempDir creates a temporary directory within tmpdir // TempDir creates a temporary directory within tmpdir
// with the name 'testname-name'. If the directory cannot // with the name 'testname-name'. If the directory cannot
// be created t.Fatal is called. // be created t.Fatal is called.
func TempDir(t *testing.T, name string) string { func TempDir(t *testing.T, name string) string {
if t != nil && t.Name() != "" { if t == nil {
name = t.Name() + "-" + name panic("argument t must be non-nil")
} }
name = t.Name() + "-" + name
name = strings.Replace(name, "/", "_", -1) name = strings.Replace(name, "/", "_", -1)
d, err := ioutil.TempDir(tmpdir, name) d, err := ioutil.TempDir(tmpdir, name)
if err != nil { if err != nil {
if t == nil {
panic(err)
}
t.Fatalf("err: %s", err) t.Fatalf("err: %s", err)
} }
t.Cleanup(func() {
if noCleanup {
t.Logf("skipping cleanup because TEST_NOCLEANUP was enabled")
return
}
os.RemoveAll(d)
})
return d return d
} }