-
Notifications
You must be signed in to change notification settings - Fork 188
/
envmock.go
129 lines (109 loc) · 2.45 KB
/
envmock.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
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
package testutil
import (
"os"
"strings"
)
// Env mocking
// MockEnvValue will store old env value, set new val. will restore old value on end.
func MockEnvValue(key, val string, fn func(nv string)) {
old := os.Getenv(key)
err := os.Setenv(key, val)
if err != nil {
panic(err)
}
fn(os.Getenv(key))
// if old is empty, unset key.
if old == "" {
err = os.Unsetenv(key)
} else {
err = os.Setenv(key, old)
}
if err != nil {
panic(err)
}
}
// MockEnvValues will store old env value, set new val. will restore old value on end.
func MockEnvValues(kvMap map[string]string, fn func()) {
backups := make(map[string]string, len(kvMap))
for key, val := range kvMap {
backups[key] = os.Getenv(key)
_ = os.Setenv(key, val)
}
fn()
for key := range kvMap {
if old := backups[key]; old == "" {
_ = os.Unsetenv(key)
} else {
_ = os.Setenv(key, old)
}
}
}
// MockOsEnvByText by env multi line text string.
// will clear all old ENV data, use given data map,
// and will recover old ENV after fn run.
//
// Usage:
//
// testutil.MockOsEnvByText(`
// APP_COMMAND = login
// APP_ENV = dev
// APP_DEBUG = true
//
// `, func() {
// // do something ...
// })
func MockOsEnvByText(envText string, fn func()) {
ss := strings.Split(envText, "\n")
mp := make(map[string]string, len(ss))
for _, line := range ss {
if line = strings.TrimSpace(line); line == "" {
continue
}
nodes := strings.SplitN(line, "=", 2)
envKey := strings.TrimSpace(nodes[0])
if len(nodes) < 2 {
mp[envKey] = ""
} else {
mp[envKey] = strings.TrimSpace(nodes[1])
}
}
MockCleanOsEnv(mp, fn)
}
// MockOsEnv by env map data. alias of MockCleanOsEnv
func MockOsEnv(mp map[string]string, fn func()) {
MockCleanOsEnv(mp, fn)
}
// backup os ENV
var envBak = os.Environ()
// ClearOSEnv info.
//
// Usage:
//
// testutil.ClearOSEnv()
// defer testutil.RevertOSEnv()
// // do something ...
func ClearOSEnv() { os.Clearenv() }
// RevertOSEnv info
func RevertOSEnv() {
os.Clearenv()
for _, str := range envBak {
nodes := strings.SplitN(str, "=", 2)
_ = os.Setenv(nodes[0], nodes[1])
}
}
// MockCleanOsEnv by env map data.
//
// will clear all old ENV data, use given data map.
// will recover old ENV after fn run.
func MockCleanOsEnv(mp map[string]string, fn func()) {
os.Clearenv()
for key, val := range mp {
_ = os.Setenv(key, val)
}
fn()
os.Clearenv()
for _, str := range envBak {
nodes := strings.SplitN(str, "=", 2)
_ = os.Setenv(nodes[0], nodes[1])
}
}