From 8f381b75cfd0a81da2aa60e2a8bc6ce2a4d79507 Mon Sep 17 00:00:00 2001
From: Tom Payne <twpayne@gmail.com>
Date: Mon, 4 May 2020 23:19:47 +0100
Subject: [PATCH] testscript: add env convenience functions

Manipulating Env.Vars is not completely trivial: you need to handle
last-entry-wins, invalid entries, and key normalization.

This commit adds Getenv, LookupEnv, and Setenv method to Env that
function identically to their counterparts in the standard os package.
---
 testscript/testscript.go      | 39 +++++++++++++++++++++++++++++++++++
 testscript/testscript_test.go | 36 ++++++++++++++++++++++++++++++++
 2 files changed, 75 insertions(+)

diff --git a/testscript/testscript.go b/testscript/testscript.go
index ea0dee0e..1fa9bde1 100644
--- a/testscript/testscript.go
+++ b/testscript/testscript.go
@@ -20,6 +20,7 @@ import (
 	"runtime"
 	"strings"
 	"sync/atomic"
+	"syscall"
 	"testing"
 	"time"
 
@@ -30,6 +31,10 @@ import (
 	"github.com/rogpeppe/go-internal/txtar"
 )
 
+// errSetenvInvalidArg is retured when Env.Setenv is called with an invalid
+// argument.
+var errSetenvInvalidArg = os.NewSyscallError("setenv", syscall.EINVAL)
+
 var execCache par.Cache
 
 // If -testwork is specified, the test prints the name of the temp directory
@@ -69,6 +74,40 @@ func (e *Env) Defer(f func()) {
 	e.ts.Defer(f)
 }
 
+// Getenv retrieves the value of the environment variable named by the key. It
+// returns the value, which will be empty if the variable is not present. To
+// distinguish between an empty value and an unset value, use LookupEnv.
+func (e *Env) Getenv(key string) string {
+	value, _ := e.LookupEnv(key)
+	return value
+}
+
+// LookupEnv retrieves the value of the environment variable named by the key.
+// If the variable is present in the environment the value (which may be empty)
+// is returned and the boolean is true. Otherwise the returned value will be
+// empty and the boolean will be false.
+func (e *Env) LookupEnv(key string) (value string, found bool) {
+	key = envvarname(key)
+	for _, v := range e.Vars {
+		pair := strings.SplitN(v, "=", 2)
+		if len(pair) == 2 && envvarname(pair[0]) == key {
+			value = pair[1]
+			found = true
+		}
+	}
+	return
+}
+
+// Setenv sets the value of the environment variable named by the key. It
+// returns an error, if any.
+func (e *Env) Setenv(key, value string) error {
+	if strings.IndexByte(key, '=') != -1 {
+		return errSetenvInvalidArg
+	}
+	e.Vars = append(e.Vars, key+"="+value)
+	return nil
+}
+
 // T returns the t argument passed to the current test by the T.Run method.
 // Note that if the tests were started by calling Run,
 // the returned value will implement testing.TB.
diff --git a/testscript/testscript_test.go b/testscript/testscript_test.go
index a5db1699..f2e7b5ab 100644
--- a/testscript/testscript_test.go
+++ b/testscript/testscript_test.go
@@ -83,6 +83,42 @@ func TestCRLFInput(t *testing.T) {
 	})
 }
 
+func TestEnv(t *testing.T) {
+	e := &Env{
+		Vars: []string{
+			"HOME=/no-home",
+			"PATH=/usr/bin",
+			"PATH=/usr/bin:/usr/local/bin",
+			"INVALID",
+		},
+	}
+
+	if got, wantHome := e.Getenv("HOME"), "/no-home"; got != wantHome {
+		t.Errorf("e.Getenv(\"HOME\") == %q, want %q", got, wantHome)
+	}
+
+	wantPath, wantPathOK := "/usr/bin:/usr/local/bin", true
+	if gotPath, gotPathOK := e.LookupEnv("PATH"); gotPath != wantPath || gotPathOK != wantPathOK {
+		t.Errorf("e.LookupEnv(\"PATH\") == %q, %v, got %q, %v", gotPath, gotPathOK, wantPath, wantPathOK)
+	}
+
+	if err := e.Setenv("HOME", "/home/user"); err != nil {
+		t.Errorf("e.Setenv(\"HOME\", \"/home/user\") == %v, want <nil>", err)
+	}
+
+	if got, wantHome := e.Getenv("HOME"), "/home/user"; got != wantHome {
+		t.Errorf("e.Getenv(\"HOME\") == %q, want %q", got, wantHome)
+	}
+
+	if gotInvalid, gotOK := e.LookupEnv("INVALID"); gotInvalid != "" || gotOK {
+		t.Errorf("e.Lookup(\"INVALID\") == %q, %v, want \"\", false", gotInvalid, gotOK)
+	}
+
+	if gotErr, wantErr := e.Setenv("invalid=key", "value"), errSetenvInvalidArg; gotErr != wantErr {
+		t.Errorf("e.Setenv(\"invalid=key\", \"value\") == %v, want %v", gotErr, wantErr)
+	}
+}
+
 func TestScripts(t *testing.T) {
 	// TODO set temp directory.
 	testDeferCount := 0