-
Notifications
You must be signed in to change notification settings - Fork 3
/
env.go
59 lines (53 loc) · 1.45 KB
/
env.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
package tf
import (
"github.com/stretchr/testify/assert"
"os"
"testing"
)
// SetEnv sets an environment variable and returns a reset function to ensure
// the environment is always returned to it's previous state:
//
// resetEnv := tf.SetEnv(t, "HOME", "/somewhere/else")
// defer resetEnv()
//
// If you would like to set multiple environment variables, see SetEnvs().
func SetEnv(t *testing.T, name, value string) (resetEnv func()) {
originalValue, exists := os.LookupEnv(name)
assert.NoError(t, os.Setenv(name, value))
return func() {
if exists {
assert.NoError(t, os.Setenv(name, originalValue))
} else {
assert.NoError(t, os.Unsetenv(name))
}
}
}
// SetEnvs works the same way as SetEnv, but on multiple environment variables:
//
// resetEnv := tf.SetEnvs(t, map[string]string{
// "HOME": "/somewhere/else",
// "DEBUG": "on",
// })
// defer resetEnv()
//
func SetEnvs(t *testing.T, env map[string]string) (resetEnv func()) {
type original struct {
value string
exists bool
}
originalValues := map[string]original{}
for name, value := range env {
originalValue, exists := os.LookupEnv(name)
originalValues[name] = original{ originalValue, exists }
assert.NoError(t, os.Setenv(name, value))
}
return func() {
for name, originalValue := range originalValues {
if originalValue.exists {
assert.NoError(t, os.Setenv(name, originalValue.value))
} else {
assert.NoError(t, os.Unsetenv(name))
}
}
}
}