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,26 +38,11 @@ import (
const FailureCode int = -1
func setEnv(key, value string) func() {
originalValue := os.Getenv(key)
os.Setenv(key, value)
if len(originalValue) > 0 {
return func() {
os.Setenv(key, originalValue)
}
func unsetEnv(t testing.TB, key string) {
if originalValue, ok := os.LookupEnv(key); ok {
t.Cleanup(func() { os.Setenv(key, originalValue) })
os.Unsetenv(key)
}
return func() {}
}
func unsetEnv(key string) func() {
originalValue := os.Getenv(key)
os.Unsetenv(key)
if len(originalValue) > 0 {
return func() {
os.Setenv(key, originalValue)
}
}
return func() {}
}
func TestHTTPProbeProxy(t *testing.T) {
@ -70,10 +55,10 @@ func TestHTTPProbeProxy(t *testing.T) {
localProxy := server.URL
defer setEnv("http_proxy", localProxy)()
defer setEnv("HTTP_PROXY", localProxy)()
defer unsetEnv("no_proxy")()
defer unsetEnv("NO_PROXY")()
t.Setenv("http_proxy", localProxy)
t.Setenv("HTTP_PROXY", localProxy)
unsetEnv(t, "no_proxy")
unsetEnv(t, "NO_PROXY")
followNonLocalRedirects := true
prober := New(followNonLocalRedirects)