From 7cbbf47f5eab450c656f2e986f9a2cbf4e40d3d7 Mon Sep 17 00:00:00 2001
From: Chris Bandy <bandy.chris@gmail.com>
Date: Sat, 15 Apr 2023 11:19:24 -0500
Subject: [PATCH] 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.
---
 pkg/probe/http/http_test.go | 31 ++++++++-----------------------
 1 file changed, 8 insertions(+), 23 deletions(-)

diff --git a/pkg/probe/http/http_test.go b/pkg/probe/http/http_test.go
index 7ae17f046b9..894d4783b12 100644
--- a/pkg/probe/http/http_test.go
+++ b/pkg/probe/http/http_test.go
@@ -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)