[tests]: Use t.Setenv to set env vars in tests (jaegertracing#4169)

## Which problem is this PR solving? ## Short description of the changes This PR replaces `os.Setenv` with `t.Setenv`. Starting from Go 1.17, we can use `t.Setenv` to set environment variable in test. The environment variable is automatically restored to its original value when the test and all its subtests complete. This ensures that each test does not start with leftover environment variables from previous completed tests. Reference: https://pkg.go.dev/testing#T.Setenv ```go func TestFoo(t *testing.T) { // before os.Setenv(key, "new value") defer os.Unsetenv(key) // after t.Setenv(key, "new value") } ``` Signed-off-by: Eng Zer Jun <[email protected]>
shubbham1215 · Feb 22, 2023 · c822ad0 · c822ad0
1 parent fc3fabc
commit c822ad0
Show file tree

Hide file tree

Showing 9 changed files with 46 additions and 91 deletions.
diff --git a/cmd/agent/app/reporter/flags_test.go b/cmd/agent/app/reporter/flags_test.go
@@ -17,7 +17,6 @@ package reporter
 import (
 	"flag"
 	"fmt"
-	"os"
 	"testing"
 
 	"github.com/spf13/cobra"
@@ -72,11 +71,8 @@ func TestBindFlags(t *testing.T) {
 	require.NoError(t, err)
 
 	b := &Options{}
-	os.Setenv("envKey1", "envVal1")
-	defer os.Unsetenv("envKey1")
-
-	os.Setenv("envKey4", "envVal4")
-	defer os.Unsetenv("envKey4")
+	t.Setenv("envKey1", "envVal1")
+	t.Setenv("envKey4", "envVal4")
 
 	b.InitFromViper(v, zap.NewNop())
 

diff --git a/cmd/flags/flags_test.go b/cmd/flags/flags_test.go
@@ -17,7 +17,6 @@ package flags
 
 import (
 	"fmt"
-	"os"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -35,11 +34,8 @@ func TestParseJaegerTags(t *testing.T) {
 		"envVar5=${envVar5:}",
 	)
 
-	os.Setenv("envKey1", "envVal1")
-	defer os.Unsetenv("envKey1")
-
-	os.Setenv("envKey4", "envVal4")
-	defer os.Unsetenv("envKey4")
+	t.Setenv("envKey1", "envVal1")
+	t.Setenv("envKey4", "envVal4")
 
 	expectedTags := map[string]string{
 		"key":     "value",

diff --git a/pkg/config/config_test.go b/pkg/config/config_test.go
@@ -18,7 +18,6 @@ package config
 import (
 	"flag"
 	"fmt"
-	"os"
 	"testing"
 	"time"
 
@@ -56,14 +55,11 @@ func TestEnv(t *testing.T) {
 	envFlag := "jaeger.test-flag"
 	actualEnvFlag := "JAEGER_TEST_FLAG"
 
-	tempEnv := os.Getenv(actualEnvFlag)
-	defer os.Setenv(actualEnvFlag, tempEnv)
-
 	addFlags := func(flagSet *flag.FlagSet) {
 		flagSet.String(envFlag, "", "")
 	}
 	expectedString := "string"
-	os.Setenv(actualEnvFlag, expectedString)
+	t.Setenv(actualEnvFlag, expectedString)
 
 	v, _ := Viperize(addFlags)
 	assert.Equal(t, expectedString, v.GetString(envFlag))

diff --git a/plugin/metrics/factory_config_test.go b/plugin/metrics/factory_config_test.go
@@ -15,27 +15,16 @@
 package metrics
 
 import (
-	"os"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
 )
 
-func clearEnv(t *testing.T) {
-	err := os.Setenv(StorageTypeEnvVar, "")
-	require.NoError(t, err)
-}
-
 func TestFactoryConfigFromEnv(t *testing.T) {
-	clearEnv(t)
-	defer clearEnv(t)
-
 	fc := FactoryConfigFromEnv()
 	assert.Empty(t, fc.MetricsStorageType)
 
-	err := os.Setenv(StorageTypeEnvVar, prometheusStorageType)
-	require.NoError(t, err)
+	t.Setenv(StorageTypeEnvVar, prometheusStorageType)
 
 	fc = FactoryConfigFromEnv()
 	assert.Equal(t, prometheusStorageType, fc.MetricsStorageType)

diff --git a/plugin/metrics/factory_test.go b/plugin/metrics/factory_test.go
@@ -97,9 +97,6 @@ func (f *configurable) InitFromViper(v *viper.Viper, logger *zap.Logger) {
 }
 
 func TestConfigurable(t *testing.T) {
-	clearEnv(t)
-	defer clearEnv(t)
-
 	f, err := NewFactory(withConfig(prometheusStorageType))
 	require.NoError(t, err)
 	assert.NotEmpty(t, f.factories)

diff --git a/plugin/sampling/strategystore/factory_config_test.go b/plugin/sampling/strategystore/factory_config_test.go
@@ -17,7 +17,6 @@ package strategystore
 
 import (
 	"io"
-	"os"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -26,122 +25,125 @@ import (
 
 func TestFactoryConfigFromEnv(t *testing.T) {
 	tests := []struct {
+		name         string
 		env          string
 		envVar       string
 		expectedType Kind
 		expectsError bool
 	}{
-		// default
 		{
+			name:         "default",
 			expectedType: Kind("file"),
 		},
-		// file on both env vars
 		{
+			name:         "file on deprecatedSamplingTypeEnvVar",
 			env:          "file",
 			envVar:       deprecatedSamplingTypeEnvVar,
 			expectedType: Kind("file"),
 		},
 		{
+			name:         "file on SamplingTypeEnvVar",
 			env:          "file",
 			envVar:       SamplingTypeEnvVar,
 			expectedType: Kind("file"),
 		},
-		// static works on the deprecated env var, but fails on the new
 		{
+			name:         "static works on the deprecatedSamplingTypeEnvVar",
 			env:          "static",
 			envVar:       deprecatedSamplingTypeEnvVar,
 			expectedType: Kind("file"),
 		},
 		{
+			name:         "static fails on the SamplingTypeEnvVar",
 			env:          "static",
 			envVar:       SamplingTypeEnvVar,
 			expectsError: true,
 		},
-		// adaptive on both env vars
 		{
+			name:         "adaptive on deprecatedSamplingTypeEnvVar",
 			env:          "adaptive",
 			envVar:       deprecatedSamplingTypeEnvVar,
 			expectedType: Kind("adaptive"),
 		},
 		{
+			name:         "adaptive on SamplingTypeEnvVar",
 			env:          "adaptive",
 			envVar:       SamplingTypeEnvVar,
 			expectedType: Kind("adaptive"),
 		},
-		// unexpected string on both env vars
 		{
+			name:         "unexpected string on deprecatedSamplingTypeEnvVar",
 			env:          "??",
 			envVar:       deprecatedSamplingTypeEnvVar,
 			expectsError: true,
 		},
 		{
+			name:         "unexpected string on SamplingTypeEnvVar",
 			env:          "??",
 			envVar:       SamplingTypeEnvVar,
 			expectsError: true,
 		},
 	}
 
 	for _, tc := range tests {
-		// clear env
-		os.Setenv(SamplingTypeEnvVar, "")
-		os.Setenv(deprecatedSamplingTypeEnvVar, "")
+		t.Run(tc.name, func(t *testing.T) {
+			if len(tc.envVar) != 0 {
+				t.Setenv(tc.envVar, tc.env)
+			}
 
-		if len(tc.envVar) != 0 {
-			err := os.Setenv(tc.envVar, tc.env)
-			require.NoError(t, err)
-		}
-
-		f, err := FactoryConfigFromEnv(io.Discard)
-		if tc.expectsError {
-			assert.Error(t, err)
-			continue
-		}
+			f, err := FactoryConfigFromEnv(io.Discard)
+			if tc.expectsError {
+				assert.Error(t, err)
+				return
+			}
 
-		require.NoError(t, err)
-		assert.Equal(t, tc.expectedType, f.StrategyStoreType)
+			require.NoError(t, err)
+			assert.Equal(t, tc.expectedType, f.StrategyStoreType)
+		})
 	}
 }
 
 func TestGetStrategyStoreTypeFromEnv(t *testing.T) {
 	tests := []struct {
+		name               string
 		deprecatedEnvValue string
 		currentEnvValue    string
 		expected           string
 	}{
-		// default to file
 		{
+			name:     "default to file",
 			expected: "file",
 		},
-		// current env var works
 		{
+			name:            "current env var works",
 			currentEnvValue: "foo",
 			expected:        "foo",
 		},
-		// current overrides deprecated
 		{
+			name:               "current overrides deprecated",
 			currentEnvValue:    "foo",
 			deprecatedEnvValue: "blerg",
 			expected:           "foo",
 		},
-		// deprecated accepted
 		{
+			name:               "deprecated accepted",
 			deprecatedEnvValue: "blerg",
 			expected:           "blerg",
 		},
-		// static is switched to file
 		{
+			name:               "static is switched to file",
 			deprecatedEnvValue: "static",
 			expected:           "file",
 		},
 	}
 
 	for _, tc := range tests {
-		err := os.Setenv(SamplingTypeEnvVar, tc.currentEnvValue)
-		require.NoError(t, err)
-		err = os.Setenv(deprecatedSamplingTypeEnvVar, tc.deprecatedEnvValue)
-		require.NoError(t, err)
+		t.Run(tc.name, func(t *testing.T) {
+			t.Setenv(SamplingTypeEnvVar, tc.currentEnvValue)
+			t.Setenv(deprecatedSamplingTypeEnvVar, tc.deprecatedEnvValue)
 
-		actual := getStrategyStoreTypeFromEnv(io.Discard)
-		assert.Equal(t, actual, tc.expected)
+			actual := getStrategyStoreTypeFromEnv(io.Discard)
+			assert.Equal(t, actual, tc.expected)
+		})
 	}
 }
diff --git a/plugin/sampling/strategystore/factory_test.go b/plugin/sampling/strategystore/factory_test.go
@@ -18,7 +18,6 @@ package strategystore
 import (
 	"errors"
 	"flag"
-	"os"
 	"testing"
 
 	"github.com/spf13/viper"
@@ -34,10 +33,6 @@ import (
 	"github.com/jaegertracing/jaeger/storage/samplingstore"
 )
 
-func clearEnv() {
-	os.Setenv(SamplingTypeEnvVar, "static")
-}
-
 var (
 	_ ss.Factory          = new(Factory)
 	_ plugin.Configurable = new(Factory)
@@ -99,8 +94,7 @@ func TestNewFactory(t *testing.T) {
 }
 
 func TestConfigurable(t *testing.T) {
-	clearEnv()
-	defer clearEnv()
+	t.Setenv(SamplingTypeEnvVar, "static")
 
 	f, err := NewFactory(FactoryConfig{StrategyStoreType: "file"})
 	require.NoError(t, err)

diff --git a/plugin/storage/factory_config_test.go b/plugin/storage/factory_config_test.go
@@ -17,32 +17,22 @@ package storage
 
 import (
 	"bytes"
-	"os"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 )
 
-func clearEnv() {
-	os.Setenv(SpanStorageTypeEnvVar, "")
-	os.Setenv(DependencyStorageTypeEnvVar, "")
-	os.Setenv(SamplingStorageTypeEnvVar, "")
-}
-
 func TestFactoryConfigFromEnv(t *testing.T) {
-	clearEnv()
-	defer clearEnv()
-
 	f := FactoryConfigFromEnvAndCLI(nil, &bytes.Buffer{})
 	assert.Equal(t, 1, len(f.SpanWriterTypes))
 	assert.Equal(t, cassandraStorageType, f.SpanWriterTypes[0])
 	assert.Equal(t, cassandraStorageType, f.SpanReaderType)
 	assert.Equal(t, cassandraStorageType, f.DependenciesStorageType)
 	assert.Empty(t, f.SamplingStorageType)
 
-	os.Setenv(SpanStorageTypeEnvVar, elasticsearchStorageType)
-	os.Setenv(DependencyStorageTypeEnvVar, memoryStorageType)
-	os.Setenv(SamplingStorageTypeEnvVar, cassandraStorageType)
+	t.Setenv(SpanStorageTypeEnvVar, elasticsearchStorageType)
+	t.Setenv(DependencyStorageTypeEnvVar, memoryStorageType)
+	t.Setenv(SamplingStorageTypeEnvVar, cassandraStorageType)
 
 	f = FactoryConfigFromEnvAndCLI(nil, &bytes.Buffer{})
 	assert.Equal(t, 1, len(f.SpanWriterTypes))
@@ -51,14 +41,14 @@ func TestFactoryConfigFromEnv(t *testing.T) {
 	assert.Equal(t, memoryStorageType, f.DependenciesStorageType)
 	assert.Equal(t, cassandraStorageType, f.SamplingStorageType)
 
-	os.Setenv(SpanStorageTypeEnvVar, elasticsearchStorageType+","+kafkaStorageType)
+	t.Setenv(SpanStorageTypeEnvVar, elasticsearchStorageType+","+kafkaStorageType)
 
 	f = FactoryConfigFromEnvAndCLI(nil, &bytes.Buffer{})
 	assert.Equal(t, 2, len(f.SpanWriterTypes))
 	assert.Equal(t, []string{elasticsearchStorageType, kafkaStorageType}, f.SpanWriterTypes)
 	assert.Equal(t, elasticsearchStorageType, f.SpanReaderType)
 
-	os.Setenv(SpanStorageTypeEnvVar, badgerStorageType)
+	t.Setenv(SpanStorageTypeEnvVar, badgerStorageType)
 
 	f = FactoryConfigFromEnvAndCLI(nil, nil)
 	assert.Equal(t, 1, len(f.SpanWriterTypes))
@@ -67,8 +57,6 @@ func TestFactoryConfigFromEnv(t *testing.T) {
 }
 
 func TestFactoryConfigFromEnvDeprecated(t *testing.T) {
-	clearEnv()
-
 	testCases := []struct {
 		args  []string
 		log   bool

diff --git a/plugin/storage/factory_test.go b/plugin/storage/factory_test.go
@@ -356,9 +356,6 @@ func (f *configurable) InitFromViper(v *viper.Viper, logger *zap.Logger) {
 }
 
 func TestConfigurable(t *testing.T) {
-	clearEnv()
-	defer clearEnv()
-
 	f, err := NewFactory(defaultCfg())
 	require.NoError(t, err)
 	assert.NotEmpty(t, f.factories)