Merge pull request kubernetes#117424 from cbandy/kubeadm-testing-setenv

Replace os.Setenv with testing.T.Setenv in tests

cmd/kubeadm/app/cmd/token_test.go

@@ -249,18 +249,14 @@ func TestNewCmdToken(t *testing.T) {
 			if _, err = f.WriteString(tc.configToWrite); err != nil {
 				t.Errorf("Unable to write test file %q: %v", fullPath, err)
 			}
-			// store the current value of the environment variable.
-			storedEnv := os.Getenv(clientcmd.RecommendedConfigPathEnvVar)
 			if tc.kubeConfigEnv != "" {
-				os.Setenv(clientcmd.RecommendedConfigPathEnvVar, tc.kubeConfigEnv)
+				t.Setenv(clientcmd.RecommendedConfigPathEnvVar, tc.kubeConfigEnv)
 			}
 			cmd.SetArgs(tc.args)
 			err := cmd.Execute()
 			if (err != nil) != tc.expectedError {
 				t.Errorf("Test case %q: newCmdToken expected error: %v, saw: %v", tc.name, tc.expectedError, (err != nil))
 			}
-			// restore the environment variable.
-			os.Setenv(clientcmd.RecommendedConfigPathEnvVar, storedEnv)
 		})
 	}
 }

cmd/kubeadm/app/preflight/checks_test.go

@@ -643,9 +643,7 @@ func TestHTTPProxyCIDRCheck(t *testing.T) {
 		},
 	}
 
-	// Save current content of *_proxy and *_PROXY variables.
-	savedEnv := resetProxyEnv(t)
-	defer restoreEnv(savedEnv)
+	resetProxyEnv(t)
 
 	for _, rt := range tests {
 		warning, _ := rt.check.Check()
@@ -725,9 +723,7 @@ func TestHTTPProxyCheck(t *testing.T) {
 		},
 	}
 
-	// Save current content of *_proxy and *_PROXY variables.
-	savedEnv := resetProxyEnv(t)
-	defer restoreEnv(savedEnv)
+	resetProxyEnv(t)
 
 	for _, rt := range tests {
 		warning, _ := rt.check.Check()
@@ -744,23 +740,19 @@ func TestHTTPProxyCheck(t *testing.T) {
 	}
 }
 
-// resetProxyEnv is helper function that unsets all *_proxy variables
-// and return previously set values as map. This can be used to restore
-// original state of the environment.
-func resetProxyEnv(t *testing.T) map[string]string {
-	savedEnv := make(map[string]string)
+// resetProxyEnv is helper function that unsets all *_proxy variables.
+func resetProxyEnv(t *testing.T) {
 	for _, e := range os.Environ() {
-		pair := strings.Split(e, "=")
-		if strings.HasSuffix(strings.ToLower(pair[0]), "_proxy") {
-			savedEnv[pair[0]] = pair[1]
-			os.Unsetenv(pair[0])
+		key, value, _ := strings.Cut(e, "=")
+		if strings.HasSuffix(strings.ToLower(key), "_proxy") {
+			t.Cleanup(func() { os.Setenv(key, value) })
+			os.Unsetenv(key)
 		}
 	}
-	t.Log("Saved environment: ", savedEnv)
 
-	os.Setenv("HTTP_PROXY", "http://proxy.example.com:3128")
-	os.Setenv("HTTPS_PROXY", "https://proxy.example.com:3128")
-	os.Setenv("NO_PROXY", "example.com,10.0.0.0/8,2001:db8::/48")
+	t.Setenv("HTTP_PROXY", "http://proxy.example.com:3128")
+	t.Setenv("HTTPS_PROXY", "https://proxy.example.com:3128")
+	t.Setenv("NO_PROXY", "example.com,10.0.0.0/8,2001:db8::/48")
 	// Check if we can reliably execute tests:
 	// ProxyFromEnvironment caches the *_proxy environment variables and
 	// if ProxyFromEnvironment already executed before our test with empty
@@ -777,15 +769,6 @@ func resetProxyEnv(t *testing.T) map[string]string {
 		t.Skip("test skipped as ProxyFromEnvironment already initialized in environment without defined HTTP proxy")
 	}
 	t.Log("http.ProxyFromEnvironment is usable, continue executing test")
-	return savedEnv
-}
-
-// restoreEnv is helper function to restores values
-// of environment variables from saved state in the map
-func restoreEnv(e map[string]string) {
-	for k, v := range e {
-		os.Setenv(k, v)
-	}
 }
 
 func TestKubeletVersionCheck(t *testing.T) {

cmd/kubeadm/test/cmd/init_test.go

@@ -27,11 +27,9 @@ import (
 	"k8s.io/apimachinery/pkg/util/version"
 )
 
-func runKubeadmInit(args ...string) (string, string, int, error) {
-	const dryRunDir = "KUBEADM_INIT_DRYRUN_DIR"
-	if err := os.Setenv(dryRunDir, os.TempDir()); err != nil {
-		panic(fmt.Sprintf("could not set the %s environment variable", dryRunDir))
-	}
+func runKubeadmInit(t testing.TB, args ...string) (string, string, int, error) {
+	t.Helper()
+	t.Setenv("KUBEADM_INIT_DRYRUN_DIR", os.TempDir())
 	kubeadmPath := getKubeadmPath()
 	kubeadmArgs := []string{"init", "--dry-run", "--ignore-preflight-errors=all"}
 	kubeadmArgs = append(kubeadmArgs, args...)
@@ -73,7 +71,7 @@ func TestCmdInitToken(t *testing.T) {
 
 	for _, rt := range initTest {
 		t.Run(rt.name, func(t *testing.T) {
-			_, _, _, err := runKubeadmInit(rt.args)
+			_, _, _, err := runKubeadmInit(t, rt.args)
 			if (err == nil) != rt.expected {
 				t.Fatalf(dedent.Dedent(`
 					CmdInitToken test case %q failed with an error: %v
@@ -112,7 +110,7 @@ func TestCmdInitKubernetesVersion(t *testing.T) {
 
 	for _, rt := range initTest {
 		t.Run(rt.name, func(t *testing.T) {
-			_, _, _, err := runKubeadmInit(rt.args)
+			_, _, _, err := runKubeadmInit(t, rt.args)
 			if (err == nil) != rt.expected {
 				t.Fatalf(dedent.Dedent(`
 					CmdInitKubernetesVersion test case %q failed with an error: %v
@@ -176,7 +174,7 @@ func TestCmdInitConfig(t *testing.T) {
 
 	for _, rt := range initTest {
 		t.Run(rt.name, func(t *testing.T) {
-			_, _, _, err := runKubeadmInit(rt.args)
+			_, _, _, err := runKubeadmInit(t, rt.args)
 			if (err == nil) != rt.expected {
 				t.Fatalf(dedent.Dedent(`
 						CmdInitConfig test case %q failed with an error: %v
@@ -225,7 +223,7 @@ func TestCmdInitAPIPort(t *testing.T) {
 
 	for _, rt := range initTest {
 		t.Run(rt.name, func(t *testing.T) {
-			_, _, _, err := runKubeadmInit(rt.args)
+			_, _, _, err := runKubeadmInit(t, rt.args)
 			if (err == nil) != rt.expected {
 				t.Fatalf(dedent.Dedent(`
 							CmdInitAPIPort test case %q failed with an error: %v
@@ -267,7 +265,7 @@ func TestCmdInitFeatureGates(t *testing.T) {
 
 	for _, rt := range initTest {
 		t.Run(rt.name, func(t *testing.T) {
-			_, _, exitcode, err := runKubeadmInit(rt.args)
+			_, _, exitcode, err := runKubeadmInit(t, rt.args)
 			if exitcode == PanicExitcode {
 				t.Fatalf(dedent.Dedent(`
 							CmdInitFeatureGates test case %q failed with an error: %v