testscript: Expose Environ() to custom commands

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

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.

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)