|
@@ -0,0 +1,92 @@
|
|
|
+/*Package env provides functions to test code that read environment variables
|
|
|
+or the current working directory.
|
|
|
+*/
|
|
|
+package env
|
|
|
+
|
|
|
+import (
|
|
|
+ "os"
|
|
|
+ "strings"
|
|
|
+
|
|
|
+ "github.com/gotestyourself/gotestyourself/assert"
|
|
|
+)
|
|
|
+
|
|
|
+type helperT interface {
|
|
|
+ Helper()
|
|
|
+}
|
|
|
+
|
|
|
+// Patch changes the value of an environment variable, and returns a
|
|
|
+// function which will reset the the value of that variable back to the
|
|
|
+// previous state.
|
|
|
+func Patch(t assert.TestingT, key, value string) func() {
|
|
|
+ if ht, ok := t.(helperT); ok {
|
|
|
+ ht.Helper()
|
|
|
+ }
|
|
|
+ oldValue, ok := os.LookupEnv(key)
|
|
|
+ assert.NilError(t, os.Setenv(key, value))
|
|
|
+ return func() {
|
|
|
+ if ht, ok := t.(helperT); ok {
|
|
|
+ ht.Helper()
|
|
|
+ }
|
|
|
+ if !ok {
|
|
|
+ assert.NilError(t, os.Unsetenv(key))
|
|
|
+ return
|
|
|
+ }
|
|
|
+ assert.NilError(t, os.Setenv(key, oldValue))
|
|
|
+ }
|
|
|
+}
|
|
|
+
|
|
|
+// PatchAll sets the environment to env, and returns a function which will
|
|
|
+// reset the environment back to the previous state.
|
|
|
+func PatchAll(t assert.TestingT, env map[string]string) func() {
|
|
|
+ if ht, ok := t.(helperT); ok {
|
|
|
+ ht.Helper()
|
|
|
+ }
|
|
|
+ oldEnv := os.Environ()
|
|
|
+ os.Clearenv()
|
|
|
+
|
|
|
+ for key, value := range env {
|
|
|
+ assert.NilError(t, os.Setenv(key, value))
|
|
|
+ }
|
|
|
+ return func() {
|
|
|
+ if ht, ok := t.(helperT); ok {
|
|
|
+ ht.Helper()
|
|
|
+ }
|
|
|
+ os.Clearenv()
|
|
|
+ for key, oldVal := range ToMap(oldEnv) {
|
|
|
+ assert.NilError(t, os.Setenv(key, oldVal))
|
|
|
+ }
|
|
|
+ }
|
|
|
+}
|
|
|
+
|
|
|
+// ToMap takes a list of strings in the format returned by os.Environ() and
|
|
|
+// returns a mapping of keys to values.
|
|
|
+func ToMap(env []string) map[string]string {
|
|
|
+ result := map[string]string{}
|
|
|
+ for _, raw := range env {
|
|
|
+ parts := strings.SplitN(raw, "=", 2)
|
|
|
+ switch len(parts) {
|
|
|
+ case 1:
|
|
|
+ result[raw] = ""
|
|
|
+ case 2:
|
|
|
+ result[parts[0]] = parts[1]
|
|
|
+ }
|
|
|
+ }
|
|
|
+ return result
|
|
|
+}
|
|
|
+
|
|
|
+// ChangeWorkingDir to the directory, and return a function which restores the
|
|
|
+// previous working directory.
|
|
|
+func ChangeWorkingDir(t assert.TestingT, dir string) func() {
|
|
|
+ if ht, ok := t.(helperT); ok {
|
|
|
+ ht.Helper()
|
|
|
+ }
|
|
|
+ cwd, err := os.Getwd()
|
|
|
+ assert.NilError(t, err)
|
|
|
+ assert.NilError(t, os.Chdir(dir))
|
|
|
+ return func() {
|
|
|
+ if ht, ok := t.(helperT); ok {
|
|
|
+ ht.Helper()
|
|
|
+ }
|
|
|
+ assert.NilError(t, os.Chdir(cwd))
|
|
|
+ }
|
|
|
+}
|