From 587957abc894d79427665377b92f7c5706de279b Mon Sep 17 00:00:00 2001
From: Paul Jolly <paul@myitcv.io>
Date: Tue, 19 Jan 2021 06:29:53 +0000
Subject: [PATCH] cmd/testscript: allow -e flag to pass env value

The -e flag for testscript currently allows the caller to specify that
environment variables are inherited from the current environment.

However this falls when the current environment does not have a value
set and you need to set a specific value. For example:

    GOBUILD=$(go env GOBUILD) testscript -e GOBUILD script.txt

Instead, allow -e to also specify a value:

    testscript -e GOBUILD=$(go env GOBUILD) script.txt

which feels more fluent/natural.
---
 cmd/testscript/README.md               |  2 +-
 cmd/testscript/help.go                 |  2 +-
 cmd/testscript/main.go                 | 16 +++++++++++----
 cmd/testscript/testdata/env_values.txt | 27 ++++++++++++++++++++++++++
 4 files changed, 41 insertions(+), 6 deletions(-)
 create mode 100644 cmd/testscript/testdata/env_values.txt

diff --git a/cmd/testscript/README.md b/cmd/testscript/README.md
index 41cc3890..3a022abe 100644
--- a/cmd/testscript/README.md
+++ b/cmd/testscript/README.md
@@ -3,7 +3,7 @@ The testscript command runs github.com/rogpeppe/go-internal/testscript scripts
 in a fresh temporary work directory tree.
 
 Usage:
-    testscript [-v] [-e VAR]... [-u] files...
+    testscript [-v] [-e VAR[=value]]... [-u] files...
 
 The testscript command is designed to make it easy to create self-contained
 reproductions of command sequences.
diff --git a/cmd/testscript/help.go b/cmd/testscript/help.go
index 2340668b..56a2e512 100644
--- a/cmd/testscript/help.go
+++ b/cmd/testscript/help.go
@@ -14,7 +14,7 @@ The testscript command runs github.com/rogpeppe/go-internal/testscript scripts
 in a fresh temporary work directory tree.
 
 Usage:
-    testscript [-v] [-e VAR]... [-u] files...
+    testscript [-v] [-e VAR[=value]]... [-u] files...
 
 The testscript command is designed to make it easy to create self-contained
 reproductions of command sequences.
diff --git a/cmd/testscript/main.go b/cmd/testscript/main.go
index 7d73d484..b96defb6 100644
--- a/cmd/testscript/main.go
+++ b/cmd/testscript/main.go
@@ -273,11 +273,19 @@ func run(runDir, filename string, update bool, verbose bool, envVars []string) e
 	if len(envVars) > 0 {
 		addSetup(func(env *testscript.Env) error {
 			for _, v := range envVars {
-				if v == "WORK" {
-					// cannot override WORK
-					continue
+				varName := v
+				if i := strings.Index(v, "="); i >= 0 {
+					varName = v[:i]
+				} else {
+					v = fmt.Sprintf("%s=%s", v, os.Getenv(v))
 				}
-				env.Vars = append(env.Vars, v+"="+os.Getenv(v))
+				switch varName {
+				case "":
+					return fmt.Errorf("invalid variable name %q", varName)
+				case "WORK":
+					return fmt.Errorf("cannot override WORK variable")
+				}
+				env.Vars = append(env.Vars, v)
 			}
 			return nil
 		})
diff --git a/cmd/testscript/testdata/env_values.txt b/cmd/testscript/testdata/env_values.txt
new file mode 100644
index 00000000..457be905
--- /dev/null
+++ b/cmd/testscript/testdata/env_values.txt
@@ -0,0 +1,27 @@
+# Test that passing env values, e.g. ENV1=val, works
+
+unquote test.txt
+
+env BLAH1=
+env BLAH2=junk
+
+# Normal operation
+testscript -e BLAH1=rubbish -e BLAH2 test.txt
+
+# It is an error to specify WORK. Note the error message
+# appears on stdout because it is written to the log output
+# of testscript, which has no concept of stderr.
+! testscript -e BLAH1=rubbish -e BLAH2 -e WORK test.txt
+stdout 'cannot override WORK variable'
+
+-- test.txt --
+>exec printf $BLAH1'\n'
+>cmp stdout stdout1.txt
+>
+>exec printf $BLAH2'\n'
+>cmp stdout stdout2.txt
+>
+>-- stdout1.txt --
+>rubbish
+>-- stdout2.txt --
+>junk