From 4f35e07b0429a9677ad257a5c18309e96c157129 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daniel=20Mart=C3=AD?= <mvdan@mvdan.cc>
Date: Mon, 24 Feb 2025 17:04:03 +0000
Subject: [PATCH] testscript: remove temp dirs when finishing once again

The recent transition from RunMain to Main meant that we started
calling os.Exit directly when Main finished.
However, we failed to spot that os.Exit would make our earlier deferred
cleanup of the temporary directory not run at all.

Spotted because my /tmp started getting filled up after a few hours
of development, and I found a suspiciously large number of
testscript-looking directories.
---
 testscript/exe.go | 80 +++++++++++++++++++++++++----------------------
 1 file changed, 42 insertions(+), 38 deletions(-)

diff --git a/testscript/exe.go b/testscript/exe.go
index 28d5698..7e6983b 100644
--- a/testscript/exe.go
+++ b/testscript/exe.go
@@ -50,50 +50,54 @@ func Main(m TestingM, commands map[string]func()) {
 	if mainf == nil {
 		// Unknown command; this is just the top-level execution of the
 		// test binary by "go test".
+		os.Exit(testingMRun(m, commands))
+	}
+	// The command being registered is being invoked, so run it, then exit.
+	os.Args[0] = cmdName
+	mainf()
+	os.Exit(0)
+}
 
-		// Set up all commands in a directory, added in $PATH.
-		tmpdir, err := os.MkdirTemp("", "testscript-main")
-		if err != nil {
-			log.Fatalf("could not set up temporary directory: %v", err)
-		}
-		defer func() {
-			if err := os.RemoveAll(tmpdir); err != nil {
-				log.Fatalf("cannot delete temporary directory: %v", err)
-			}
-		}()
-		bindir := filepath.Join(tmpdir, "bin")
-		if err := os.MkdirAll(bindir, 0o777); err != nil {
-			log.Fatalf("could not set up PATH binary directory: %v", err)
+// testingMRun exists just so that we can use `defer`, given that [Main] above uses [os.Exit].
+func testingMRun(m TestingM, commands map[string]func()) int {
+	// Set up all commands in a directory, added in $PATH.
+	tmpdir, err := os.MkdirTemp("", "testscript-main")
+	if err != nil {
+		log.Fatalf("could not set up temporary directory: %v", err)
+	}
+	defer func() {
+		if err := os.RemoveAll(tmpdir); err != nil {
+			log.Fatalf("cannot delete temporary directory: %v", err)
 		}
-		os.Setenv("PATH", bindir+string(filepath.ListSeparator)+os.Getenv("PATH"))
+	}()
+	bindir := filepath.Join(tmpdir, "bin")
+	if err := os.MkdirAll(bindir, 0o777); err != nil {
+		log.Fatalf("could not set up PATH binary directory: %v", err)
+	}
+	os.Setenv("PATH", bindir+string(filepath.ListSeparator)+os.Getenv("PATH"))
 
-		// We're not in a subcommand.
-		for name := range commands {
-			// Set up this command in the directory we added to $PATH.
-			binfile := filepath.Join(bindir, name)
-			if runtime.GOOS == "windows" {
-				binfile += ".exe"
-			}
-			binpath, err := os.Executable()
-			if err == nil {
-				err = copyBinary(binpath, binfile)
-			}
-			if err != nil {
-				log.Fatalf("could not set up %s in $PATH: %v", name, err)
-			}
-			scriptCmds[name] = func(ts *TestScript, neg bool, args []string) {
-				if ts.params.RequireExplicitExec {
-					ts.Fatalf("use 'exec %s' rather than '%s' (because RequireExplicitExec is enabled)", name, name)
-				}
-				ts.cmdExec(neg, append([]string{name}, args...))
+	// We're not in a subcommand.
+	for name := range commands {
+		// Set up this command in the directory we added to $PATH.
+		binfile := filepath.Join(bindir, name)
+		if runtime.GOOS == "windows" {
+			binfile += ".exe"
+		}
+		binpath, err := os.Executable()
+		if err == nil {
+			err = copyBinary(binpath, binfile)
+		}
+		if err != nil {
+			log.Fatalf("could not set up %s in $PATH: %v", name, err)
+		}
+		scriptCmds[name] = func(ts *TestScript, neg bool, args []string) {
+			if ts.params.RequireExplicitExec {
+				ts.Fatalf("use 'exec %s' rather than '%s' (because RequireExplicitExec is enabled)", name, name)
 			}
+			ts.cmdExec(neg, append([]string{name}, args...))
 		}
-		os.Exit(m.Run())
 	}
-	// The command being registered is being invoked, so run it, then exit.
-	os.Args[0] = cmdName
-	mainf()
-	os.Exit(0)
+	return m.Run()
 }
 
 // Deprecated: use [Main], as the only reason for returning exit codes