Replace os.Setenv by testing.Setenv in testhelper.go

internal/config/config_test.go

@@ -14,7 +14,7 @@ import (
 )
 
 func TestConfigApplyGlobalState(t *testing.T) {
-	t.Cleanup(testhelper.TempEnv(map[string]string{"SSL_CERT_DIR": "unmodified"}))
+	testhelper.TempEnv(t, map[string]string{"SSL_CERT_DIR": "unmodified"})
 
 	config := &Config{SslCertDir: ""}
 	config.ApplyGlobalState()

internal/executable/executable_test.go

@@ -52,8 +52,7 @@ func TestNewSuccess(t *testing.T) {
 
 	for _, tc := range testCases {
 		t.Run(tc.desc, func(t *testing.T) {
-			restoreEnv := testhelper.TempEnv(tc.environment)
-			defer restoreEnv()
+			testhelper.TempEnv(t, tc.environment)
 
 			fake := tc.fakeOs
 			fake.Setup()
@@ -89,8 +88,7 @@ func TestNewFailure(t *testing.T) {
 
 	for _, tc := range testCases {
 		t.Run(tc.desc, func(t *testing.T) {
-			restoreEnv := testhelper.TempEnv(tc.environment)
-			defer restoreEnv()
+			testhelper.TempEnv(t, tc.environment)
 
 			fake := tc.fakeOs
 			fake.Setup()

internal/sshenv/sshenv_test.go

@@ -32,8 +32,7 @@ func TestNewFromEnv(t *testing.T) {
 
 	for _, tc := range tests {
 		t.Run(tc.desc, func(t *testing.T) {
-			restoreEnv := testhelper.TempEnv(tc.environment)
-			defer restoreEnv()
+			testhelper.TempEnv(t, tc.environment)
 
 			require.Equal(t, NewFromEnv(), tc.want)
 		})
@@ -41,9 +40,7 @@ func TestNewFromEnv(t *testing.T) {
 }
 
 func TestRemoteAddrFromEnv(t *testing.T) {
-	cleanup, err := testhelper.Setenv(SSHConnectionEnv, "127.0.0.1 0")
-	require.NoError(t, err)
-	defer cleanup()
+	t.Setenv(SSHConnectionEnv, "127.0.0.1 0")
 
 	require.Equal(t, remoteAddrFromEnv(), "127.0.0.1")
 }

internal/testhelper/testhelper.go

@@ -11,17 +11,9 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func TempEnv(env map[string]string) func() {
-	var original = make(map[string]string)
+func TempEnv(t *testing.T, env map[string]string) {
 	for key, value := range env {
-		original[key] = os.Getenv(key)
-		os.Setenv(key, value)
-	}
-
-	return func() {
-		for key, originalValue := range original {
-			os.Setenv(key, originalValue)
-		}
+		t.Setenv(key, value)
 	}
 }
 
@@ -62,9 +54,3 @@ func getTestDataDir() (string, error) {
 
 	return path.Join(path.Dir(currentFile), "testdata"), nil
 }
-
-func Setenv(key, value string) (func(), error) {
-	oldValue := os.Getenv(key)
-	err := os.Setenv(key, value)
-	return func() { os.Setenv(key, oldValue) }, err
-}