From 73a01f60031968389a55313c63b6ee9250ba3226 Mon Sep 17 00:00:00 2001
From: James Telfer <792299+jamestelfer@users.noreply.github.com>
Date: Wed, 11 Dec 2024 10:52:46 +1100
Subject: [PATCH] test: use unsetEnv helper to manage global state

---
 src/plugin/config_test.go | 26 +++++++++++++++++---------
 1 file changed, 17 insertions(+), 9 deletions(-)

diff --git a/src/plugin/config_test.go b/src/plugin/config_test.go
index 6ef80ef..5754738 100644
--- a/src/plugin/config_test.go
+++ b/src/plugin/config_test.go
@@ -10,11 +10,11 @@ import (
 )
 
 func TestFailOnMissingEnvironment(t *testing.T) {
-	unsetEnvironmentVariables()
-
 	var config plugin.Config
 	fetcher := plugin.EnvironmentConfigFetcher{}
 
+	unsetEnv(t, "BUILDKITE_PLUGIN_EXAMPLE_GO_MESSAGE")
+
 	err := fetcher.Fetch(&config)
 
 	expectedErr := "required key BUILDKITE_PLUGIN_EXAMPLE_GO_MESSAGE missing value"
@@ -22,9 +22,6 @@ func TestFailOnMissingEnvironment(t *testing.T) {
 }
 
 func TestFetchConfigFromEnvironment(t *testing.T) {
-	unsetEnvironmentVariables()
-	defer unsetEnvironmentVariables()
-
 	var config plugin.Config
 	fetcher := plugin.EnvironmentConfigFetcher{}
 
@@ -36,8 +33,19 @@ func TestFetchConfigFromEnvironment(t *testing.T) {
 	assert.Equal(t, "test-message", config.Message, "fetched message should match environment")
 }
 
-// Unsets environment variables through an all-in-one function. Extend this with additional environment variables as
-// needed.
-func unsetEnvironmentVariables() {
-	os.Unsetenv("BUILDKITE_PLUGIN_EXAMPLE_GO_MESSAGE")
+func unsetEnv(t *testing.T, key string) {
+	t.Helper()
+
+	// ensure state is restored correctly
+	currValue, exists := os.LookupEnv(key)
+	t.Cleanup(func() {
+		if exists {
+			os.Setenv(key, currValue)
+		} else {
+			os.Unsetenv(key)
+		}
+	})
+
+	// clear the value
+	os.Unsetenv(key)
 }