cmd/testscript: add support to pass through env vars with -e (#54)

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] files...
+    testscript [-v] [-e VAR]... files...
 
 The testscript command is designed to make it easy to create self-contained
 reproductions of command sequences.
@@ -31,6 +31,11 @@ proxy server. See the documentation for
 github.com/rogpeppe/go-internal/goproxytest for details on the format of these
 files/directories.
 
+Environment variables can be passed through to each script with the -e flag,
+where VAR is the name of the variable. Variables override testscript-defined
+values, with the exception of WORK which cannot be overridden. The -e flag can
+appear multiple times to specify multiple variables.
+
 Examples
 ========
 

cmd/testscript/main.go

@@ -28,6 +28,19 @@ const (
 	goModProxyDir = ".gomodproxy"
 )
 
+type envVarsFlag struct {
+	vals []string
+}
+
+func (e *envVarsFlag) String() string {
+	return fmt.Sprintf("%v", e.vals)
+}
+
+func (e *envVarsFlag) Set(v string) error {
+	e.vals = append(e.vals, v)
+	return nil
+}
+
 func main() {
 	os.Exit(main1())
 }
@@ -49,8 +62,10 @@ func mainerr() (retErr error) {
 	fs.Usage = func() {
 		mainUsage(os.Stderr)
 	}
+	var envVars envVarsFlag
 	fWork := fs.Bool("work", false, "print temporary work directory and do not remove when done")
 	fVerbose := fs.Bool("v", false, "run tests verbosely")
+	fs.Var(&envVars, "e", "pass through environment variable to script (can appear multiple times)")
 	if err := fs.Parse(os.Args[1:]); err != nil {
 		return err
 	}
@@ -77,7 +92,7 @@ func mainerr() (retErr error) {
 		if err := os.Mkdir(runDir, 0777); err != nil {
 			return fmt.Errorf("failed to create a run directory within %v for %v: %v", td, fileName, err)
 		}
-		if err := run(runDir, fileName, *fVerbose); err != nil {
+		if err := run(runDir, fileName, *fVerbose, envVars.vals); err != nil {
 			return err
 		}
 	}
@@ -127,7 +142,7 @@ func (r runner) Verbose() bool {
 	return r.verbose
 }
 
-func run(runDir, fileName string, verbose bool) error {
+func run(runDir, fileName string, verbose bool, envVars []string) error {
 	var ar *txtar.Archive
 	var err error
 
@@ -184,26 +199,44 @@ func run(runDir, fileName string, verbose bool) error {
 		}
 	}
 
-	if len(gomodProxy.Files) > 0 {
-		srv, err := goproxytest.NewServer(mods, "")
-		if err != nil {
-			return fmt.Errorf("cannot start proxy for %v: %v", fileName, err)
-		}
-		defer srv.Close()
-
+	addSetup := func(f func(env *testscript.Env) error) {
 		origSetup := p.Setup
-
 		p.Setup = func(env *testscript.Env) error {
 			if origSetup != nil {
 				if err := origSetup(env); err != nil {
 					return err
 				}
 			}
+			return f(env)
+		}
+	}
+
+	if len(gomodProxy.Files) > 0 {
+		srv, err := goproxytest.NewServer(mods, "")
+		if err != nil {
+			return fmt.Errorf("cannot start proxy for %v: %v", fileName, err)
+		}
+		defer srv.Close()
+
+		addSetup(func(env *testscript.Env) error {
 			// Add GOPROXY after calling the original setup
 			// so that it overrides any GOPROXY set there.
 			env.Vars = append(env.Vars, "GOPROXY="+srv.URL)
 			return nil
-		}
+		})
+	}
+
+	if len(envVars) > 0 {
+		addSetup(func(env *testscript.Env) error {
+			for _, v := range envVars {
+				if v == "WORK" {
+					// cannot override WORK
+					continue
+				}
+				env.Vars = append(env.Vars, v+"="+os.Getenv(v))
+			}
+			return nil
+		})
 	}
 
 	r := runner{

cmd/testscript/main_test.go

@@ -10,10 +10,12 @@ import (
 	"os"
 	"os/exec"
 	"path/filepath"
+	"runtime"
 	"strings"
 	"testing"
 
 	"github.com/rogpeppe/go-internal/gotooltest"
+	"github.com/rogpeppe/go-internal/internal/os/execpath"
 	"github.com/rogpeppe/go-internal/testscript"
 )
 
@@ -24,6 +26,10 @@ func TestMain(m *testing.M) {
 }
 
 func TestScripts(t *testing.T) {
+	if _, err := exec.LookPath("go"); err != nil {
+		t.Fatalf("need go in PATH for these tests")
+	}
+
 	var stderr bytes.Buffer
 	cmd := exec.Command("go", "env", "GOMOD")
 	cmd.Stderr = &stderr
@@ -46,7 +52,9 @@ func TestScripts(t *testing.T) {
 			return nil
 		},
 		Cmds: map[string]func(ts *testscript.TestScript, neg bool, args []string){
-			"unquote": unquote,
+			"unquote":        unquote,
+			"dropgofrompath": dropgofrompath,
+			"setfilegoproxy": setfilegoproxy,
 		},
 	}
 	if err := gotooltest.Setup(&p); err != nil {
@@ -69,3 +77,36 @@ func unquote(ts *testscript.TestScript, neg bool, args []string) {
 		ts.Check(err)
 	}
 }
+
+func dropgofrompath(ts *testscript.TestScript, neg bool, args []string) {
+	if neg {
+		ts.Fatalf("unsupported: ! dropgofrompath")
+	}
+	var newPath []string
+	for _, d := range filepath.SplitList(ts.Getenv("PATH")) {
+		getenv := func(k string) string {
+			if k == "PATH" {
+				return d
+			}
+			return ts.Getenv(k)
+		}
+		if _, err := execpath.Look("go", getenv); err != nil {
+			newPath = append(newPath, d)
+		}
+	}
+	ts.Setenv("PATH", strings.Join(newPath, string(filepath.ListSeparator)))
+}
+
+func setfilegoproxy(ts *testscript.TestScript, neg bool, args []string) {
+	if neg {
+		ts.Fatalf("unsupported: ! setfilegoproxy")
+	}
+	path := args[0]
+	path = filepath.ToSlash(path)
+	// probably sufficient to just handle spaces
+	path = strings.Replace(path, " ", "%20", -1)
+	if runtime.GOOS == "windows" {
+		path = "/" + path
+	}
+	ts.Setenv("GOPROXY", "file://"+path)
+}

cmd/testscript/testdata/env_var_no_go.txt

@@ -0,0 +1,100 @@
+# Test passing environment variables to scripts with no go command on PATH
+#
+# This is the counterpart to env_var_with_go.txt
+
+unquote noproxy.txt
+unquote withproxy.txt
+dropgofrompath
+
+env BANANA=banana
+env GOPATH=$WORK/ourgopath
+
+env GOPROXY=
+
+# no GOPROXY, no pass-through, no proxy
+testscript -v noproxy.txt
+stdout ^BANANA=$
+stdout ^GOPATH=$
+stdout ^GOPROXY=$
+! stderr .+
+
+# no GOPROXY, no pass-through, with proxy
+testscript -v withproxy.txt
+stdout ^BANANA=$
+stdout ^GOPATH=$
+stdout ^GOPROXY=http://.*/mod$
+! stderr .+
+
+# no GOPROXY, with pass-through, no proxy
+testscript -v -e BANANA -e GOPATH -e GOPROXY noproxy.txt
+stdout ^BANANA=banana$
+stdout ^GOPATH=${WORK@R}[/\\]ourgopath$
+stdout ^GOPROXY=$
+! stderr .+
+
+# no GOPROXY, with pass-through, with proxy
+testscript -v -e BANANA -e GOPATH -e GOPROXY withproxy.txt
+stdout ^BANANA=banana$
+stdout ^GOPATH=${WORK@R}[/\\]ourgopath$
+stdout ^GOPROXY=$
+! stderr .+
+
+setfilegoproxy $WORK/proxy
+
+# with GOPROXY, no pass-through, no proxy
+testscript -v noproxy.txt
+stdout ^BANANA=$
+stdout ^GOPATH=$
+stdout ^GOPROXY=$
+! stderr .+
+
+# with GOPROXY, no pass-through, with proxy
+testscript -v withproxy.txt
+stdout ^BANANA=$
+stdout ^GOPATH=$
+stdout ^GOPROXY=http://.*/mod$
+! stderr .+
+
+# with GOPROXY, with pass-through, no proxy
+testscript -v -e BANANA -e GOPATH -e GOPROXY noproxy.txt
+stdout ^BANANA=banana$
+stdout ^GOPATH=${WORK@R}[/\\]ourgopath$
+stdout ^GOPROXY=$GOPROXY$
+! stderr .+
+
+# with GOPROXY, with pass-through, with proxy
+testscript -v -e BANANA -e GOPATH -e GOPROXY withproxy.txt
+stdout ^BANANA=banana$
+stdout ^GOPATH=${WORK@R}[/\\]ourgopath$
+stdout ^GOPROXY=$GOPROXY$
+! stderr .+
+
+-- noproxy.txt --
+>env BANANA
+>env GOPATH
+>env GOPROXY
+
+-- withproxy.txt --
+>env BANANA
+>env GOPATH
+>env GOPROXY
+
+>-- .gomodproxy/fruit.com_v1.0.0/.mod --
+>module fruit.com
+>
+>-- .gomodproxy/fruit.com_v1.0.0/.info --
+>{"Version":"v1.0.0","Time":"2018-10-22T18:45:39Z"}
+>
+>-- .gomodproxy/fruit.com_v1.0.0/go.mod --
+>module fruit.com
+>
+>-- .gomodproxy/fruit.com_v1.0.0/fruit/fruit.go --
+>package fruit
+>
+>const Apple = "apple"
+>-- .gomodproxy/fruit.com_v1.0.0/coretest/coretest.go --
+>// package coretest becomes a candidate for the missing
+>// core import in main above
+>package coretest
+>
+>const Mandarin = "mandarin"