Skip to content

Commit

Permalink
[tests]: Use t.Setenv to set env vars in tests (#4169)
Browse files Browse the repository at this point in the history
## 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 <engzerjun@gmail.com>
  • Loading branch information
Juneezee committed Jan 19, 2023
1 parent 0237ca4 commit c7df877
Show file tree
Hide file tree
Showing 9 changed files with 46 additions and 91 deletions.
8 changes: 2 additions & 6 deletions cmd/agent/app/reporter/flags_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@ package reporter
import (
"flag"
"fmt"
"os"
"testing"

"github.com/spf13/cobra"
Expand Down Expand Up @@ -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())

Expand Down
8 changes: 2 additions & 6 deletions cmd/flags/flags_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@ package flags

import (
"fmt"
"os"
"testing"

"github.com/stretchr/testify/assert"
Expand All @@ -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",
Expand Down
6 changes: 1 addition & 5 deletions pkg/config/config_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,6 @@ package config
import (
"flag"
"fmt"
"os"
"testing"
"time"

Expand Down Expand Up @@ -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))
Expand Down
13 changes: 1 addition & 12 deletions plugin/metrics/factory_config_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
3 changes: 0 additions & 3 deletions plugin/metrics/factory_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
66 changes: 34 additions & 32 deletions plugin/sampling/strategystore/factory_config_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@ package strategystore

import (
"io"
"os"
"testing"

"github.com/stretchr/testify/assert"
Expand All @@ -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)
})
}
}
8 changes: 1 addition & 7 deletions plugin/sampling/strategystore/factory_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,6 @@ package strategystore
import (
"errors"
"flag"
"os"
"testing"

"github.com/spf13/viper"
Expand All @@ -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)
Expand Down Expand Up @@ -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)
Expand Down
22 changes: 5 additions & 17 deletions plugin/storage/factory_config_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand All @@ -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))
Expand All @@ -67,8 +57,6 @@ func TestFactoryConfigFromEnv(t *testing.T) {
}

func TestFactoryConfigFromEnvDeprecated(t *testing.T) {
clearEnv()

testCases := []struct {
args []string
log bool
Expand Down
3 changes: 0 additions & 3 deletions plugin/storage/factory_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down

0 comments on commit c7df877

Please sign in to comment.