testing: add TB.Setenv

Add a new method TB.Setenv that'll set environment variables
only for the isolated lifetime of the test, and will clean up
and unset these variables when the test ends.
This method disables the test or benchmark from running in
parallel.

Fixes #41260

Change-Id: I0a18f094ec1c6ec3157b4b12993ea3075e2e9867
GitHub-Last-Rev: 0ca12fa565318f350b927e2ef94f3b4f792c75c2
GitHub-Pull-Request: golang/go#41857
Reviewed-on: https://go-review.googlesource.com/c/go/+/260577
Trust: Daniel Martí <mvdan@mvdan.cc>
Trust: Emmanuel Odeke <emmanuel@orijtech.com>
Run-TryBot: Daniel Martí <mvdan@mvdan.cc>
TryBot-Result: Go Bot <gobot@golang.org>
Reviewed-by: roger peppe <rogpeppe@gmail.com>

1 files changed, 83 insertions, 0 deletions

diff --git a/src/testing/testing_test.go b/src/testing/testing_test.go
index 0f096980ca..55a4df4739 100644
--- a/src/testing/testing_test.go
+++ b/src/testing/testing_test.go
@@ -109,3 +109,86 @@ func testTempDir(t *testing.T) {
 		t.Errorf("unexpected %d files in TempDir: %v", len(files), files)
 	}
 }
+
+func TestSetenv(t *testing.T) {
+	tests := []struct {
+		name               string
+		key                string
+		initialValueExists bool
+		initialValue       string
+		newValue           string
+	}{
+		{
+			name:               "initial value exists",
+			key:                "GO_TEST_KEY_1",
+			initialValueExists: true,
+			initialValue:       "111",
+			newValue:           "222",
+		},
+		{
+			name:               "initial value exists but empty",
+			key:                "GO_TEST_KEY_2",
+			initialValueExists: true,
+			initialValue:       "",
+			newValue:           "222",
+		},
+		{
+			name:               "initial value is not exists",
+			key:                "GO_TEST_KEY_3",
+			initialValueExists: false,
+			initialValue:       "",
+			newValue:           "222",
+		},
+	}
+
+	for _, test := range tests {
+		if test.initialValueExists {
+			if err := os.Setenv(test.key, test.initialValue); err != nil {
+				t.Fatalf("unable to set env: got %v", err)
+			}
+		} else {
+			os.Unsetenv(test.key)
+		}
+
+		t.Run(test.name, func(t *testing.T) {
+			t.Setenv(test.key, test.newValue)
+			if os.Getenv(test.key) != test.newValue {
+				t.Fatalf("unexpected value after t.Setenv: got %s, want %s", os.Getenv(test.key), test.newValue)
+			}
+		})
+
+		got, exists := os.LookupEnv(test.key)
+		if got != test.initialValue {
+			t.Fatalf("unexpected value after t.Setenv cleanup: got %s, want %s", got, test.initialValue)
+		}
+		if exists != test.initialValueExists {
+			t.Fatalf("unexpected value after t.Setenv cleanup: got %t, want %t", exists, test.initialValueExists)
+		}
+	}
+}
+
+func TestSetenvWithParallelAfterSetenv(t *testing.T) {
+	defer func() {
+		want := "testing: t.Parallel called after t.Setenv; cannot set environment variables in parallel tests"
+		if got := recover(); got != want {
+			t.Fatalf("expected panic; got %#v want %q", got, want)
+		}
+	}()
+
+	t.Setenv("GO_TEST_KEY_1", "value")
+
+	t.Parallel()
+}
+
+func TestSetenvWithParallelBeforeSetenv(t *testing.T) {
+	defer func() {
+		want := "testing: t.Setenv called after t.Parallel; cannot set environment variables in parallel tests"
+		if got := recover(); got != want {
+			t.Fatalf("expected panic; got %#v want %q", got, want)
+		}
+	}()
+
+	t.Parallel()
+
+	t.Setenv("GO_TEST_KEY_1", "value")
+}