test: Add testutil.CreateTemp, remove testutil.TempDir

mafredri · mafredri · commit 635cadc5700b · 2023-01-17T11:35:54.000Z
diff --git a/cli/server_test.go b/cli/server_test.go
@@ -1134,7 +1134,7 @@ func TestServer(t *testing.T) {
 
 			random, err := cryptorand.String(5)
 			require.NoError(t, err)
-			tmpdir := testutil.TempDir(t)
+			tmpdir := t.TempDir()
 			fiName := fmt.Sprint(tmpdir, "/coder-logging-test-", random)
 
 			root, _ := clitest.New(t,
@@ -1163,9 +1163,7 @@ func TestServer(t *testing.T) {
 			ctx, cancelFunc := context.WithCancel(context.Background())
 			defer cancelFunc()
 
-			tmpdir := testutil.TempDir(t)
-			fi, err := os.CreateTemp(tmpdir, "coder-logging-test-*")
-			require.NoError(t, err)
+			fi := testutil.CreateTemp(t, "", "coder-logging-test-*")
 
 			root, _ := clitest.New(t,
 				"server",
@@ -1193,9 +1191,7 @@ func TestServer(t *testing.T) {
 			ctx, cancelFunc := context.WithCancel(context.Background())
 			defer cancelFunc()
 
-			tmpdir := testutil.TempDir(t)
-			fi, err := os.CreateTemp(tmpdir, "coder-logging-test-*")
-			require.NoError(t, err)
+			fi := testutil.CreateTemp(t, "", "coder-logging-test-*")
 
 			root, _ := clitest.New(t,
 				"server",
@@ -1223,9 +1219,7 @@ func TestServer(t *testing.T) {
 			ctx, cancelFunc := context.WithCancel(context.Background())
 			defer cancelFunc()
 
-			tmpdir := testutil.TempDir(t)
-			fi, err := os.CreateTemp(tmpdir, "coder-logging-test-*")
-			require.NoError(t, err)
+			fi := testutil.CreateTemp(t, "", "coder-logging-test-*")
 
 			root, _ := clitest.New(t,
 				"server",
@@ -1266,13 +1260,9 @@ func TestServer(t *testing.T) {
 			ctx, cancelFunc := context.WithCancel(context.Background())
 			defer cancelFunc()
 
-			tmpdir := testutil.TempDir(t)
-			fi1, err := os.CreateTemp(tmpdir, "coder-logging-test-*")
-			require.NoError(t, err)
-			fi2, err := os.CreateTemp(tmpdir, "coder-logging-test-*")
-			require.NoError(t, err)
-			fi3, err := os.CreateTemp(tmpdir, "coder-logging-test-*")
-			require.NoError(t, err)
+			fi1 := testutil.CreateTemp(t, "", "coder-logging-test-*")
+			fi2 := testutil.CreateTemp(t, "", "coder-logging-test-*")
+			fi3 := testutil.CreateTemp(t, "", "coder-logging-test-*")
 
 			// NOTE(mafredri): This test might end up downloading Terraform
 			// which can take a long time and end up failing the test.
diff --git a/testutil/temp.go b/testutil/temp.go
@@ -1,75 +1,32 @@
 package testutil
 
 import (
-	"io/fs"
 	"os"
-	"path/filepath"
-	"runtime"
 	"testing"
-	"time"
+
+	"github.com/stretchr/testify/require"
 )
 
-func TempDir(t *testing.T) string {
+// CreateTemp is a convenience function for creating a temporary file, like
+// os.CreateTemp, but it also registers a cleanup function to close and remove
+// the file.
+func CreateTemp(t *testing.T, dir, pattern string) *os.File {
 	t.Helper()
-	dir := t.TempDir()
-
-	// On Windows, os.RemoveAll can fail with (among other):
-	//
-	//	The process cannot access the file because it is being used by another process.
-	//
-	// We try to work around this issue by removing files and directories in
-	// a loop until os.RemoveAll succeeds or we give up.
-	if runtime.GOOS == "windows" {
-		t.Cleanup(func() {
-			var dirs, files []string
-			_ = filepath.WalkDir(dir, func(path string, d fs.DirEntry, err error) error {
-				if err != nil {
-					return err
-				}
-				if d.IsDir() && path != dir {
-					dirs = append(dirs, path)
-					return nil
-				}
-				files = append(files, path)
-				return nil
-			})
-			removeFilesAndDirs := func() {
-				// Try to remove all files so that the directories are empty.
-				var newFiles []string
-				for _, path := range files {
-					if err := os.Remove(path); err != nil {
-						newFiles = append(newFiles, path)
-					}
-				}
-				files = newFiles
-
-				// Remove directories in reverse order (~depth first).
-				var newDirs []string
-				for i := len(dirs) - 1; i >= 0; i-- {
-					if err := os.Remove(dirs[i]); err != nil {
-						newDirs = append([]string{dirs[i]}, newDirs...)
-					}
-				}
-				dirs = newDirs
-			}
 
-			var (
-				start   = time.Now()
-				timeout = 5 * time.Second
-				err     error
-			)
-			for time.Since(start) < timeout {
-				removeFilesAndDirs()
-				err = os.RemoveAll(dir)
-				if err == nil {
-					return
-				}
-				time.Sleep(100 * time.Millisecond)
-			}
-			// Note that even though we're giving up here, t.TempDir
-			// cleanup _could_ still succeed (or fail the test).
-			t.Logf("TempDir: delete %q: giving up after %s: %v", dir, timeout, err)
-		})
+	if dir == "" {
+		dir = t.TempDir()
 	}
-	return dir
+	f, err := os.CreateTemp(dir, pattern)
+	require.NoError(t, err, "create temp file")
+	t.Cleanup(func() {
+		err = f.Close()
+		if err != nil {
+			t.Logf("CreateTemp: Cleanup: close failed for %q: %v", f.Name(), err)
+		}
+		err = os.Remove(f.Name())
+		if err != nil {
+			t.Logf("CreateTemp: Cleanup: remove failed for %q: %v", f.Name(), err)
+		}
+	})
+	return f
 }
