From fa4df28fcd5d8523b56c95820f9655af87b0ccd0 Mon Sep 17 00:00:00 2001
From: Mahmood Ali <mahmood@hashicorp.com>
Date: Wed, 13 Oct 2021 17:26:56 -0400
Subject: [PATCH] tests: ensure that tests restore env-var values (#11309)

Fix a test corruption issue, where a test accidentally unsets
the `NOMAD_LICENSE` environment variable, that's relied on by some
tests.

As a habit, tests should always restore the environment variable value
on test completion. Golang 1.17 introduced
[`t.Setenv`](https://pkg.go.dev/testing#T.Setenv) to address this issue.
However, as 1.0.x and 1.1.x branches target golang 1.15 and 1.16, I
opted to use a helper function to ease backports.
---
 command/helpers_test.go |  3 +--
 command/meta_test.go    | 29 +++++++++--------------------
 command/util_test.go    | 15 +++++++++++++++
 3 files changed, 25 insertions(+), 22 deletions(-)

diff --git a/command/helpers_test.go b/command/helpers_test.go
index b24e80b74..c5e54e853 100644
--- a/command/helpers_test.go
+++ b/command/helpers_test.go
@@ -346,8 +346,7 @@ job "example" {
 }
 `
 
-	os.Setenv("NOMAD_VAR_var4", "from-envvar")
-	defer os.Unsetenv("NOMAD_VAR_var4")
+	setEnv(t, "NOMAD_VAR_var4", "from-envvar")
 
 	cliArgs := []string{`var2=from-cli`}
 	fileVars := `var3 = "from-varfile"`
diff --git a/command/meta_test.go b/command/meta_test.go
index 34a97d8d6..e8ba4209a 100644
--- a/command/meta_test.go
+++ b/command/meta_test.go
@@ -5,12 +5,11 @@ import (
 	"os"
 	"reflect"
 	"sort"
-	"strings"
 	"testing"
 
 	"github.com/kr/pty"
 	"github.com/mitchellh/cli"
-	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestMeta_FlagSet(t *testing.T) {
@@ -90,7 +89,7 @@ func TestMeta_Colorize(t *testing.T) {
 		{
 			Name: "disable colors via env var",
 			SetupFn: func(t *testing.T, m *Meta) {
-				os.Setenv(EnvNomadCLINoColor, "1")
+				setEnv(t, EnvNomadCLINoColor, "1")
 				m.SetupUi([]string{})
 			},
 			ExpectColor: false,
@@ -105,7 +104,7 @@ func TestMeta_Colorize(t *testing.T) {
 		{
 			Name: "force colors via env var",
 			SetupFn: func(t *testing.T, m *Meta) {
-				os.Setenv(EnvNomadCLIForceColor, "1")
+				setEnv(t, EnvNomadCLIForceColor, "1")
 				m.SetupUi([]string{})
 			},
 			ExpectColor: true,
@@ -120,7 +119,7 @@ func TestMeta_Colorize(t *testing.T) {
 		{
 			Name: "no color take predecence over force color via env var",
 			SetupFn: func(t *testing.T, m *Meta) {
-				os.Setenv(EnvNomadCLINoColor, "1")
+				setEnv(t, EnvNomadCLINoColor, "1")
 				m.SetupUi([]string{"-force-color"})
 			},
 			ExpectColor: false,
@@ -131,22 +130,16 @@ func TestMeta_Colorize(t *testing.T) {
 		t.Run(tc.Name, func(t *testing.T) {
 			// Create fake test terminal.
 			_, tty, err := pty.Open()
-			if err != nil {
-				t.Fatalf("%v", err)
-			}
+			require.NoError(t, err)
 			defer tty.Close()
 
 			oldStdout := os.Stdout
 			defer func() { os.Stdout = oldStdout }()
 			os.Stdout = tty
 
-			// Make sure Nomad environment variables are clean.
-			for _, envVar := range os.Environ() {
-				if strings.HasPrefix(envVar, "NOMAD") {
-					k := strings.SplitN(envVar, "=", 2)[0]
-					os.Unsetenv(k)
-				}
-			}
+			// Make sure color related environment variables are clean.
+			setEnv(t, EnvNomadCLIForceColor, "")
+			setEnv(t, EnvNomadCLINoColor, "")
 
 			// Run test case.
 			m := &Meta{}
@@ -154,11 +147,7 @@ func TestMeta_Colorize(t *testing.T) {
 				tc.SetupFn(t, m)
 			}
 
-			if tc.ExpectColor {
-				assert.False(t, m.Colorize().Disable)
-			} else {
-				assert.True(t, m.Colorize().Disable)
-			}
+			require.Equal(t, !tc.ExpectColor, m.Colorize().Disable)
 		})
 	}
 }
diff --git a/command/util_test.go b/command/util_test.go
index d1fd96a45..11ea7cd36 100644
--- a/command/util_test.go
+++ b/command/util_test.go
@@ -1,6 +1,7 @@
 package command
 
 import (
+	"os"
 	"testing"
 
 	"github.com/hashicorp/nomad/api"
@@ -106,3 +107,17 @@ func testMultiRegionJob(jobID, region, datacenter string) *api.Job {
 
 	return job
 }
+
+// setEnv wraps os.Setenv(key, value) and restores the environment variable to initial value in test cleanup
+func setEnv(t *testing.T, key, value string) {
+	initial, ok := os.LookupEnv(key)
+	os.Setenv(key, value)
+
+	t.Cleanup(func() {
+		if ok {
+			os.Setenv(key, initial)
+		} else {
+			os.Unsetenv(key)
+		}
+	})
+}