You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
119 lines
2.3 KiB
Go
119 lines
2.3 KiB
Go
package mock
|
|
|
|
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 text string.
|
|
// will clear all old ENV data, use given data map.
|
|
// will recover old ENV after fn run.
|
|
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])
|
|
}
|
|
}
|