From e17c166b17b27f0ab94cc2ecc97fa0979a6e98e2 Mon Sep 17 00:00:00 2001 From: Eng Zer Jun Date: Thu, 19 Jan 2023 22:17:41 +0800 Subject: [PATCH] test: use `t.Setenv` to set env vars in tests 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 --- cmd/agent/app/reporter/flags_test.go | 8 +-- cmd/flags/flags_test.go | 8 +-- pkg/config/config_test.go | 6 +- plugin/metrics/factory_config_test.go | 13 +--- plugin/metrics/factory_test.go | 3 - .../strategystore/factory_config_test.go | 66 ++++++++++--------- plugin/sampling/strategystore/factory_test.go | 8 +-- plugin/storage/factory_config_test.go | 22 ++----- plugin/storage/factory_test.go | 3 - 9 files changed, 46 insertions(+), 91 deletions(-) diff --git a/cmd/agent/app/reporter/flags_test.go b/cmd/agent/app/reporter/flags_test.go index 1baa00f2603..f024fb3ca04 100644 --- 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 index 7aceefc8a7c..c272ac094db 100644 --- 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 index df53017e9b9..fac9ceade85 100644 --- 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 index 82c9ad6aab1..c6caeaebd10 100644 --- 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 index 13a9444beb6..b24ba5d64a4 100644 --- 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 index 035b144e66d..5cd02e075e9 100644 --- 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,55 +25,60 @@ 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, @@ -82,66 +86,64 @@ func TestFactoryConfigFromEnv(t *testing.T) { } 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 index 1bf2886ffbe..66b2eca3c64 100644 --- 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 index cfce7ca22ab..fd737abe7dd 100644 --- a/plugin/storage/factory_config_test.go +++ b/plugin/storage/factory_config_test.go @@ -17,22 +17,12 @@ 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]) @@ -40,9 +30,9 @@ func TestFactoryConfigFromEnv(t *testing.T) { 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 index 66e6c086d38..0dc23ad1983 100644 --- 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)