From cf1945d0c3ba94ec983cc79a3baae6bc95d5ead5 Mon Sep 17 00:00:00 2001
From: Eng Zer Jun <engzerjun@gmail.com>
Date: Mon, 5 Jun 2023 17:31:03 +0800
Subject: [PATCH] Tests: use `t.Setenv` to set env vars (#69516)

This commit replaces `os.Setenv` with `t.Setenv` in tests. The
environment variable is automatically restored to its original value
when the test and all its subtests complete.

Reference: https://pkg.go.dev/testing#T.Setenv

Signed-off-by: Eng Zer Jun <engzerjun@gmail.com>
---
 pkg/build/config/genmetadata_test.go          |  6 ++--
 pkg/build/env/fallback_test.go                |  5 +--
 .../commands/diagnostics_test.go              | 18 +++-------
 pkg/infra/tracing/tracing_test.go             |  8 +----
 .../notification_policy_types_test.go         | 10 ++----
 .../dashboards/config_reader_test.go          |  4 +--
 .../datasources/config_reader_test.go         |  4 +--
 .../notifiers/config_reader_test.go           |  4 +--
 .../plugins/config_reader_test.go             |  7 +---
 .../provisioning/values/values_test.go        | 23 +++----------
 pkg/setting/dynamic_settings_test.go          |  8 +----
 pkg/setting/expanders_test.go                 |  3 +-
 pkg/setting/setting_test.go                   | 33 ++++++++-----------
 13 files changed, 35 insertions(+), 98 deletions(-)

diff --git a/pkg/build/config/genmetadata_test.go b/pkg/build/config/genmetadata_test.go
index 52dfca9096a..78413e2b7aa 100644
--- a/pkg/build/config/genmetadata_test.go
+++ b/pkg/build/config/genmetadata_test.go
@@ -72,10 +72,8 @@ func setUpEnv(t *testing.T, envMap map[string]string) {
 	t.Helper()
 
 	os.Clearenv()
-	err := os.Setenv("DRONE_COMMIT", "abcd12345")
-	require.NoError(t, err)
+	t.Setenv("DRONE_COMMIT", "abcd12345")
 	for k, v := range envMap {
-		err := os.Setenv(k, v)
-		require.NoError(t, err)
+		t.Setenv(k, v)
 	}
 }
diff --git a/pkg/build/env/fallback_test.go b/pkg/build/env/fallback_test.go
index d14e13ee774..2940a33201f 100644
--- a/pkg/build/env/fallback_test.go
+++ b/pkg/build/env/fallback_test.go
@@ -135,9 +135,6 @@ func setEnv(t *testing.T, key, value string) string {
 	t.Helper()
 	os.Clearenv()
 
-	err := os.Setenv(key, value)
-	if err != nil {
-		require.NoError(t, err)
-	}
+	t.Setenv(key, value)
 	return key
 }
diff --git a/pkg/cmd/grafana-server/commands/diagnostics_test.go b/pkg/cmd/grafana-server/commands/diagnostics_test.go
index 6f99a7f7b8b..0cd59f6a4e9 100644
--- a/pkg/cmd/grafana-server/commands/diagnostics_test.go
+++ b/pkg/cmd/grafana-server/commands/diagnostics_test.go
@@ -2,7 +2,6 @@ package commands
 
 import (
 	"fmt"
-	"os"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -25,18 +24,14 @@ func TestProfilingDiagnostics(t *testing.T) {
 
 	for i, tc := range tcs {
 		t.Run(fmt.Sprintf("testcase %d", i), func(t *testing.T) {
-			os.Clearenv()
 			if tc.enabledEnv != "" {
-				err := os.Setenv(profilingEnabledEnvName, tc.enabledEnv)
-				assert.NoError(t, err)
+				t.Setenv(profilingEnabledEnvName, tc.enabledEnv)
 			}
 			if tc.addrEnv != "" {
-				err := os.Setenv(profilingAddrEnvName, tc.addrEnv)
-				assert.NoError(t, err)
+				t.Setenv(profilingAddrEnvName, tc.addrEnv)
 			}
 			if tc.portEnv != "" {
-				err := os.Setenv(profilingPortEnvName, tc.portEnv)
-				assert.NoError(t, err)
+				t.Setenv(profilingPortEnvName, tc.portEnv)
 			}
 			err := tc.defaults.overrideWithEnv()
 			assert.NoError(t, err)
@@ -61,14 +56,11 @@ func TestTracingDiagnostics(t *testing.T) {
 
 	for i, tc := range tcs {
 		t.Run(fmt.Sprintf("testcase %d", i), func(t *testing.T) {
-			os.Clearenv()
 			if tc.enabledEnv != "" {
-				err := os.Setenv(tracingEnabledEnvName, tc.enabledEnv)
-				assert.NoError(t, err)
+				t.Setenv(tracingEnabledEnvName, tc.enabledEnv)
 			}
 			if tc.fileEnv != "" {
-				err := os.Setenv(tracingFileEnvName, tc.fileEnv)
-				assert.NoError(t, err)
+				t.Setenv(tracingFileEnvName, tc.fileEnv)
 			}
 			err := tc.defaults.overrideWithEnv()
 			assert.NoError(t, err)
diff --git a/pkg/infra/tracing/tracing_test.go b/pkg/infra/tracing/tracing_test.go
index 17bc34bcc72..9ed2ce1027b 100644
--- a/pkg/infra/tracing/tracing_test.go
+++ b/pkg/infra/tracing/tracing_test.go
@@ -1,7 +1,6 @@
 package tracing
 
 import (
-	"os"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -141,13 +140,8 @@ func TestTracingConfig(t *testing.T) {
 			// export envioronment variables
 			if test.Env != nil {
 				for k, v := range test.Env {
-					assert.NoError(t, os.Setenv(k, v))
+					t.Setenv(k, v)
 				}
-				defer func() {
-					for k := range test.Env {
-						assert.NoError(t, os.Unsetenv(k))
-					}
-				}()
 			}
 			// parse config sections
 			cfg := setting.NewCfg()
diff --git a/pkg/services/provisioning/alerting/notification_policy_types_test.go b/pkg/services/provisioning/alerting/notification_policy_types_test.go
index 15080a154f3..5ec98c8722f 100644
--- a/pkg/services/provisioning/alerting/notification_policy_types_test.go
+++ b/pkg/services/provisioning/alerting/notification_policy_types_test.go
@@ -1,7 +1,6 @@
 package alerting
 
 import (
-	"os"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -13,11 +12,8 @@ func TestNotificationPolicy(t *testing.T) {
 		envKey   = "NOTIFIER_EMAIL_REMINDER_FREQUENCY"
 		envValue = "4h"
 	)
-	err := os.Setenv(envKey, envValue)
-	require.NoError(t, err)
-	defer func() {
-		_ = os.Unsetenv(envKey)
-	}()
+	t.Setenv(envKey, envValue)
+
 	data := `orgId: 123
 receiver: test
 continue: true
@@ -25,7 +21,7 @@ repeat_interval: ${NOTIFIER_EMAIL_REMINDER_FREQUENCY}
 `
 	var model NotificiationPolicyV1
 
-	err = yaml.Unmarshal([]byte(data), &model)
+	err := yaml.Unmarshal([]byte(data), &model)
 	require.NoError(t, err)
 	np, err := model.mapToModel()
 	require.NoError(t, err)
diff --git a/pkg/services/provisioning/dashboards/config_reader_test.go b/pkg/services/provisioning/dashboards/config_reader_test.go
index ded54885ed4..07f0ade89a1 100644
--- a/pkg/services/provisioning/dashboards/config_reader_test.go
+++ b/pkg/services/provisioning/dashboards/config_reader_test.go
@@ -3,7 +3,6 @@ package dashboards
 import (
 	"context"
 	"errors"
-	"os"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -47,10 +46,9 @@ func TestDashboardsAsConfig(t *testing.T) {
 		})
 
 		t.Run("Can read config file version 1 format", func(t *testing.T) {
-			_ = os.Setenv("TEST_VAR", "general")
+			t.Setenv("TEST_VAR", "general")
 			cfgProvider := configReader{path: simpleDashboardConfig, log: logger, orgService: orgFake}
 			cfg, err := cfgProvider.readConfig(context.Background())
-			_ = os.Unsetenv("TEST_VAR")
 			require.NoError(t, err)
 
 			validateDashboardAsConfig(t, cfg)
diff --git a/pkg/services/provisioning/datasources/config_reader_test.go b/pkg/services/provisioning/datasources/config_reader_test.go
index 6de16a7fb63..42efa29b3c1 100644
--- a/pkg/services/provisioning/datasources/config_reader_test.go
+++ b/pkg/services/provisioning/datasources/config_reader_test.go
@@ -2,7 +2,6 @@ package datasources
 
 import (
 	"context"
-	"os"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -192,10 +191,9 @@ func TestDatasourceAsConfig(t *testing.T) {
 	})
 
 	t.Run("can read all properties from version 1", func(t *testing.T) {
-		_ = os.Setenv("TEST_VAR", "name")
+		t.Setenv("TEST_VAR", "name")
 		cfgProvider := &configReader{log: log.New("test logger"), orgService: &orgtest.FakeOrgService{}}
 		cfg, err := cfgProvider.readConfig(context.Background(), allProperties)
-		_ = os.Unsetenv("TEST_VAR")
 		if err != nil {
 			t.Fatalf("readConfig return an error %v", err)
 		}
diff --git a/pkg/services/provisioning/notifiers/config_reader_test.go b/pkg/services/provisioning/notifiers/config_reader_test.go
index 80eb7d785d1..f9772cc340e 100644
--- a/pkg/services/provisioning/notifiers/config_reader_test.go
+++ b/pkg/services/provisioning/notifiers/config_reader_test.go
@@ -3,7 +3,6 @@ package notifiers
 import (
 	"context"
 	"fmt"
-	"os"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -70,7 +69,7 @@ func TestNotificationAsConfig(t *testing.T) {
 
 		t.Run("Can read correct properties", func(t *testing.T) {
 			setup()
-			_ = os.Setenv("TEST_VAR", "default")
+			t.Setenv("TEST_VAR", "default")
 			cfgProvider := &configReader{
 				orgService:        orgService,
 				encryptionService: encryptionService,
@@ -78,7 +77,6 @@ func TestNotificationAsConfig(t *testing.T) {
 			}
 
 			cfg, err := cfgProvider.readConfig(context.Background(), correctProperties)
-			_ = os.Unsetenv("TEST_VAR")
 			if err != nil {
 				t.Fatalf("readConfig return an error %v", err)
 			}
diff --git a/pkg/services/provisioning/plugins/config_reader_test.go b/pkg/services/provisioning/plugins/config_reader_test.go
index 3440c5848d7..a1b936f7b72 100644
--- a/pkg/services/provisioning/plugins/config_reader_test.go
+++ b/pkg/services/provisioning/plugins/config_reader_test.go
@@ -2,7 +2,6 @@ package plugins
 
 import (
 	"context"
-	"os"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -55,11 +54,7 @@ func TestConfigReader(t *testing.T) {
 			},
 		}
 
-		err := os.Setenv("ENABLE_PLUGIN_VAR", "test-plugin")
-		require.NoError(t, err)
-		t.Cleanup(func() {
-			_ = os.Unsetenv("ENABLE_PLUGIN_VAR")
-		})
+		t.Setenv("ENABLE_PLUGIN_VAR", "test-plugin")
 
 		cfgProvider := newConfigReader(log.New("test logger"), pm)
 		cfg, err := cfgProvider.readConfig(context.Background(), correctProperties)
diff --git a/pkg/services/provisioning/values/values_test.go b/pkg/services/provisioning/values/values_test.go
index 70574280b3b..18152f8ead4 100644
--- a/pkg/services/provisioning/values/values_test.go
+++ b/pkg/services/provisioning/values/values_test.go
@@ -16,25 +16,10 @@ import (
 
 func TestValues(t *testing.T) {
 	t.Run("Values", func(t *testing.T) {
-		err := os.Setenv("INT", "1")
-		require.NoError(t, err)
-		err = os.Setenv("STRING", "test")
-		require.NoError(t, err)
-		err = os.Setenv("EMPTYSTRING", "")
-		require.NoError(t, err)
-		err = os.Setenv("BOOL", "true")
-		require.NoError(t, err)
-
-		defer func() {
-			err := os.Unsetenv("INT")
-			require.NoError(t, err)
-			err = os.Unsetenv("STRING")
-			require.NoError(t, err)
-			err = os.Unsetenv("EMPTYSTRING")
-			require.NoError(t, err)
-			err = os.Unsetenv("BOOL")
-			require.NoError(t, err)
-		}()
+		t.Setenv("INT", "1")
+		t.Setenv("STRING", "test")
+		t.Setenv("EMPTYSTRING", "")
+		t.Setenv("BOOL", "true")
 
 		t.Run("IntValue", func(t *testing.T) {
 			type Data struct {
diff --git a/pkg/setting/dynamic_settings_test.go b/pkg/setting/dynamic_settings_test.go
index 8b18415d592..4f827bcb99d 100644
--- a/pkg/setting/dynamic_settings_test.go
+++ b/pkg/setting/dynamic_settings_test.go
@@ -1,7 +1,6 @@
 package setting
 
 import (
-	"os"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -14,12 +13,7 @@ func TestDynamicSettingsSupport_Override(t *testing.T) {
 	keyName := "bar"
 	expected := "dynamic value"
 
-	err := os.Setenv(envKey, expected)
-	require.NoError(t, err)
-	defer func() {
-		err := os.Unsetenv(envKey)
-		require.NoError(t, err)
-	}()
+	t.Setenv(envKey, expected)
 
 	value := cfg.SectionWithEnvOverrides(sectionName).Key(keyName).MustString("default value")
 	require.Equal(t, expected, value)
diff --git a/pkg/setting/expanders_test.go b/pkg/setting/expanders_test.go
index fbf2b6bea64..909701079db 100644
--- a/pkg/setting/expanders_test.go
+++ b/pkg/setting/expanders_test.go
@@ -14,8 +14,7 @@ import (
 func TestExpandVar_EnvSuccessful(t *testing.T) {
 	const key = "GF_TEST_SETTING_EXPANDER_ENV"
 	const expected = "aurora borealis"
-	err := os.Setenv(key, expected)
-	require.NoError(t, err)
+	t.Setenv(key, expected)
 
 	// expanded format
 	{
diff --git a/pkg/setting/setting_test.go b/pkg/setting/setting_test.go
index ac93ceff817..b1bd27b32bf 100644
--- a/pkg/setting/setting_test.go
+++ b/pkg/setting/setting_test.go
@@ -64,11 +64,10 @@ func TestLoadingSettings(t *testing.T) {
 	})
 
 	t.Run("Should be able to override via environment variables", func(t *testing.T) {
-		err := os.Setenv("GF_SECURITY_ADMIN_USER", "superduper")
-		require.NoError(t, err)
+		t.Setenv("GF_SECURITY_ADMIN_USER", "superduper")
 
 		cfg := NewCfg()
-		err = cfg.Load(CommandLineArgs{HomePath: "../../"})
+		err := cfg.Load(CommandLineArgs{HomePath: "../../"})
 		require.Nil(t, err)
 
 		require.Equal(t, "superduper", cfg.AdminUser)
@@ -77,22 +76,20 @@ func TestLoadingSettings(t *testing.T) {
 	})
 
 	t.Run("Should replace password when defined in environment", func(t *testing.T) {
-		err := os.Setenv("GF_SECURITY_ADMIN_PASSWORD", "supersecret")
-		require.NoError(t, err)
+		t.Setenv("GF_SECURITY_ADMIN_PASSWORD", "supersecret")
 
 		cfg := NewCfg()
-		err = cfg.Load(CommandLineArgs{HomePath: "../../"})
+		err := cfg.Load(CommandLineArgs{HomePath: "../../"})
 		require.Nil(t, err)
 
 		require.Contains(t, appliedEnvOverrides, "GF_SECURITY_ADMIN_PASSWORD=*********")
 	})
 
 	t.Run("Should replace password in URL when url environment is defined", func(t *testing.T) {
-		err := os.Setenv("GF_DATABASE_URL", "mysql://user:secret@localhost:3306/database")
-		require.NoError(t, err)
+		t.Setenv("GF_DATABASE_URL", "mysql://user:secret@localhost:3306/database")
 
 		cfg := NewCfg()
-		err = cfg.Load(CommandLineArgs{HomePath: "../../"})
+		err := cfg.Load(CommandLineArgs{HomePath: "../../"})
 		require.Nil(t, err)
 
 		require.Contains(t, appliedEnvOverrides, "GF_DATABASE_URL=mysql://user:xxxxx@localhost:3306/database")
@@ -208,10 +205,9 @@ func TestLoadingSettings(t *testing.T) {
 
 	t.Run("Can use environment variables in config values", func(t *testing.T) {
 		if runtime.GOOS == windows {
-			err := os.Setenv("GF_DATA_PATH", `c:\tmp\env_override`)
-			require.NoError(t, err)
+			t.Setenv("GF_DATA_PATH", `c:\tmp\env_override`)
 			cfg := NewCfg()
-			err = cfg.Load(CommandLineArgs{
+			err := cfg.Load(CommandLineArgs{
 				HomePath: "../../",
 				Args:     []string{"cfg:paths.data=${GF_DATA_PATH}"},
 			})
@@ -219,10 +215,9 @@ func TestLoadingSettings(t *testing.T) {
 
 			require.Equal(t, `c:\tmp\env_override`, cfg.DataPath)
 		} else {
-			err := os.Setenv("GF_DATA_PATH", "/tmp/env_override")
-			require.NoError(t, err)
+			t.Setenv("GF_DATA_PATH", "/tmp/env_override")
 			cfg := NewCfg()
-			err = cfg.Load(CommandLineArgs{
+			err := cfg.Load(CommandLineArgs{
 				HomePath: "../../",
 				Args:     []string{"cfg:paths.data=${GF_DATA_PATH}"},
 			})
@@ -282,12 +277,10 @@ func TestLoadingSettings(t *testing.T) {
 	})
 
 	t.Run("grafana.com API URL can be set separately from grafana.com URL", func(t *testing.T) {
-		err := os.Setenv("GF_GRAFANA_NET_URL", "https://grafana-dev.com")
-		require.NoError(t, err)
-		err = os.Setenv("GF_GRAFANA_COM_API_URL", "http://grafana-dev.internal/api")
-		require.NoError(t, err)
+		t.Setenv("GF_GRAFANA_NET_URL", "https://grafana-dev.com")
+		t.Setenv("GF_GRAFANA_COM_API_URL", "http://grafana-dev.internal/api")
 		cfg := NewCfg()
-		err = cfg.Load(CommandLineArgs{HomePath: "../../", Config: "../../conf/defaults.ini"})
+		err := cfg.Load(CommandLineArgs{HomePath: "../../", Config: "../../conf/defaults.ini"})
 		require.Nil(t, err)
 		require.Equal(t, "https://grafana-dev.com", cfg.GrafanaComURL)
 		require.Equal(t, "http://grafana-dev.internal/api", cfg.GrafanaComAPIURL)