From 675fe71657d96849683534c7552f584cea561364 Mon Sep 17 00:00:00 2001
From: Abhinav Gupta <mail@abhinavg.net>
Date: Sat, 11 May 2024 14:39:09 -0700
Subject: [PATCH] testscript: Expose Environ() to custom commands

This adds an Environ method to TestScript
that looks and acts similarly to os.Environ.

This gives further visibility into the test script state
to custom command implementations,
allowing them to implement functionality that is otherwise limited
only to the built-in commands.

For example, combining this and MkAbs, a custom command can use os/exec
to run a command with the same environment as the test script,
with different behavior than the built-in TestScript.Exec method.
---
 testscript/testdata/environ.txt |  8 ++++++++
 testscript/testscript.go        | 12 ++++++++++--
 testscript/testscript_test.go   | 13 +++++++++++++
 3 files changed, 31 insertions(+), 2 deletions(-)
 create mode 100644 testscript/testdata/environ.txt

diff --git a/testscript/testdata/environ.txt b/testscript/testdata/environ.txt
new file mode 100644
index 00000000..8096ce7a
--- /dev/null
+++ b/testscript/testdata/environ.txt
@@ -0,0 +1,8 @@
+env FOO_A=1 FOO_B=2 FOO_C=3
+printEnvPrefix FOO_
+cmp stdout want.txt
+
+-- want.txt --
+FOO_A=1
+FOO_B=2
+FOO_C=3
diff --git a/testscript/testscript.go b/testscript/testscript.go
index a6b3e5c8..17568181 100644
--- a/testscript/testscript.go
+++ b/testscript/testscript.go
@@ -959,7 +959,7 @@ func (ts *TestScript) exec(command string, args ...string) (stdout, stderr strin
 		return "", "", err
 	}
 	cmd.Dir = ts.cd
-	cmd.Env = append(ts.env, "PWD="+ts.cd)
+	cmd.Env = append(ts.Environ(), "PWD="+ts.cd)
 	cmd.Stdin = strings.NewReader(ts.stdin)
 	var stdoutBuf, stderrBuf strings.Builder
 	cmd.Stdout = &stdoutBuf
@@ -1012,7 +1012,7 @@ func (ts *TestScript) execBackground(command string, args ...string) (*exec.Cmd,
 		return nil, err
 	}
 	cmd.Dir = ts.cd
-	cmd.Env = append(ts.env, "PWD="+ts.cd)
+	cmd.Env = append(ts.Environ(), "PWD="+ts.cd)
 	var stdoutBuf, stderrBuf strings.Builder
 	cmd.Stdin = strings.NewReader(ts.stdin)
 	cmd.Stdout = &stdoutBuf
@@ -1222,6 +1222,14 @@ func (ts *TestScript) Getenv(key string) string {
 	return ts.envMap[envvarname(key)]
 }
 
+// Environ returns a copy of environment variables set on the TestScript
+// as a series of key=value strings.
+func (ts *TestScript) Environ() []string {
+	env := make([]string, len(ts.env))
+	copy(env, ts.env)
+	return env
+}
+
 // parse parses a single line as a list of space-separated arguments
 // subject to environment variable expansion (but not resplitting).
 // Single quotes around text disable splitting and expansion.
diff --git a/testscript/testscript_test.go b/testscript/testscript_test.go
index f01b468f..97a1abf6 100644
--- a/testscript/testscript_test.go
+++ b/testscript/testscript_test.go
@@ -284,6 +284,19 @@ func TestScripts(t *testing.T) {
 					ts.Fatalf("cannot chdir: %v", err)
 				}
 			},
+			"printEnvPrefix": func(ts *TestScript, neg bool, args []string) {
+				if neg || len(args) != 1 {
+					ts.Fatalf("usage: printEnvPrefix <prefix>")
+				}
+
+				w := ts.Stdout()
+				prefix := args[0]
+				for _, kv := range ts.Environ() {
+					if strings.HasPrefix(kv, prefix) {
+						fmt.Fprintln(w, kv)
+					}
+				}
+			},
 		},
 		Setup: func(env *Env) error {
 			infos, err := os.ReadDir(env.WorkDir)