Cleanup environment in tests that call os.Unsetenv

testing.T.Cleanup ensures the environment is restored after a test and
any of its parallel sub-tests. It's possible that these can be
simplified further to T.Setenv(key, ""), but I did not investigate.
This commit is contained in:
Chris Bandy 2023-04-15 11:19:24 -05:00
parent d38ac7e7c6
commit 7cbbf47f5e

View File

@ -38,27 +38,12 @@ import (
const FailureCode int = -1 const FailureCode int = -1
func setEnv(key, value string) func() { func unsetEnv(t testing.TB, key string) {
originalValue := os.Getenv(key) if originalValue, ok := os.LookupEnv(key); ok {
os.Setenv(key, value) t.Cleanup(func() { os.Setenv(key, originalValue) })
if len(originalValue) > 0 {
return func() {
os.Setenv(key, originalValue)
}
}
return func() {}
}
func unsetEnv(key string) func() {
originalValue := os.Getenv(key)
os.Unsetenv(key) os.Unsetenv(key)
if len(originalValue) > 0 {
return func() {
os.Setenv(key, originalValue)
} }
} }
return func() {}
}
func TestHTTPProbeProxy(t *testing.T) { func TestHTTPProbeProxy(t *testing.T) {
res := "welcome to http probe proxy" res := "welcome to http probe proxy"
@ -70,10 +55,10 @@ func TestHTTPProbeProxy(t *testing.T) {
localProxy := server.URL localProxy := server.URL
defer setEnv("http_proxy", localProxy)() t.Setenv("http_proxy", localProxy)
defer setEnv("HTTP_PROXY", localProxy)() t.Setenv("HTTP_PROXY", localProxy)
defer unsetEnv("no_proxy")() unsetEnv(t, "no_proxy")
defer unsetEnv("NO_PROXY")() unsetEnv(t, "NO_PROXY")
followNonLocalRedirects := true followNonLocalRedirects := true
prober := New(followNonLocalRedirects) prober := New(followNonLocalRedirects)