Skip to content

Commit b27b34c

Browse files
JuneezeeJuneezee
authored
test: use T.Setenv to set env vars in tests (#739)
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 <[email protected]> Signed-off-by: Eng Zer Jun <[email protected]>
1 parent bd3d277 commit b27b34c

File tree

2 files changed

+60
-72
lines changed

2 files changed

+60
-72
lines changed

‎pkg/config/config_internal_test.go‎

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@ package config
1515

1616
import (
1717
"flag"
18-
"os"
1918
"strconv"
2019
"testing"
2120

@@ -37,7 +36,7 @@ func TestGetEnv(t *testing.T) {
3736
var successVal = "success"
3837
var failVal = "failure"
3938

40-
os.Setenv(key, successVal)
39+
t.Setenv(key, successVal)
4140

4241
result := getEnv(key+"bla", failVal)
4342
h.Equals(t, failVal, result)
@@ -51,7 +50,7 @@ func TestGetIntEnv(t *testing.T) {
5150
var successVal = 1
5251
var failVal = 0
5352

54-
os.Setenv(key, strconv.Itoa(successVal))
53+
t.Setenv(key, strconv.Itoa(successVal))
5554

5655
result := getIntEnv(key+"bla", failVal)
5756
h.Equals(t, failVal, result)
@@ -64,7 +63,7 @@ func TestGetIntEnv(t *testing.T) {
6463
t.Errorf("getIntEnv did not panic")
6564
}
6665
}()
67-
os.Setenv(key, "hi")
66+
t.Setenv(key, "hi")
6867
getIntEnv(key, 0)
6968
}
7069

@@ -73,7 +72,7 @@ func TestGetBoolEnv(t *testing.T) {
7372
var successVal = true
7473
var failVal = false
7574

76-
os.Setenv(key, strconv.FormatBool(successVal))
75+
t.Setenv(key, strconv.FormatBool(successVal))
7776

7877
result := getBoolEnv(key+"bla", failVal)
7978
h.Equals(t, failVal, result)
@@ -86,7 +85,7 @@ func TestGetBoolEnv(t *testing.T) {
8685
t.Errorf("getBoolEnv did not panic")
8786
}
8887
}()
89-
os.Setenv(key, "hi")
88+
t.Setenv(key, "hi")
9089
getBoolEnv(key, false)
9190
}
9291

@@ -99,7 +98,7 @@ func TestIsConfigProvided(t *testing.T) {
9998
result = isConfigProvided(cliArgName, envVarName)
10099
h.Equals(t, true, result)
101100

102-
os.Setenv(envVarName, value)
101+
t.Setenv(envVarName, value)
103102
result = isConfigProvided(cliArgName, envVarName)
104103
h.Equals(t, true, result)
105104
}

‎pkg/config/config_test.go‎

Lines changed: 54 additions & 65 deletions
Original file line numberDiff line numberDiff line change
@@ -25,44 +25,33 @@ import (
2525
"github.com/rs/zerolog/log"
2626
)
2727

28-
var envForTest = map[string]string{}
29-
3028
func resetFlagsForTest() {
3129
flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ContinueOnError)
3230
os.Args = []string{"cmd"}
33-
for key := range envForTest {
34-
os.Unsetenv(key)
35-
}
36-
}
37-
38-
func setEnvForTest(key string, val string) {
39-
os.Setenv(key, val)
40-
envForTest[key] = val
4131
}
4232

4333
func TestParseCliArgsEnvSuccess(t *testing.T) {
44-
resetFlagsForTest()
45-
setEnvForTest("USE_PROVIDER_ID", "true")
46-
setEnvForTest("DELETE_LOCAL_DATA", "false")
47-
setEnvForTest("DRY_RUN", "true")
48-
setEnvForTest("ENABLE_SCHEDULED_EVENT_DRAINING", "true")
49-
setEnvForTest("ENABLE_SPOT_INTERRUPTION_DRAINING", "false")
50-
setEnvForTest("ENABLE_SQS_TERMINATION_DRAINING", "false")
51-
setEnvForTest("ENABLE_REBALANCE_MONITORING", "true")
52-
setEnvForTest("ENABLE_REBALANCE_DRAINING", "true")
53-
setEnvForTest("GRACE_PERIOD", "12345")
54-
setEnvForTest("IGNORE_DAEMON_SETS", "false")
55-
setEnvForTest("KUBERNETES_SERVICE_HOST", "KUBERNETES_SERVICE_HOST")
56-
setEnvForTest("KUBERNETES_SERVICE_PORT", "KUBERNETES_SERVICE_PORT")
57-
setEnvForTest("NODE_NAME", "NODE_NAME")
58-
setEnvForTest("NODE_TERMINATION_GRACE_PERIOD", "12345")
59-
setEnvForTest("INSTANCE_METADATA_URL", "INSTANCE_METADATA_URL")
60-
setEnvForTest("POD_TERMINATION_GRACE_PERIOD", "12345")
61-
setEnvForTest("WEBHOOK_URL", "WEBHOOK_URL")
62-
setEnvForTest("WEBHOOK_HEADERS", "WEBHOOK_HEADERS")
63-
setEnvForTest("WEBHOOK_TEMPLATE", "WEBHOOK_TEMPLATE")
64-
setEnvForTest("METADATA_TRIES", "100")
65-
setEnvForTest("CORDON_ONLY", "false")
34+
t.Setenv("USE_PROVIDER_ID", "true")
35+
t.Setenv("DELETE_LOCAL_DATA", "false")
36+
t.Setenv("DRY_RUN", "true")
37+
t.Setenv("ENABLE_SCHEDULED_EVENT_DRAINING", "true")
38+
t.Setenv("ENABLE_SPOT_INTERRUPTION_DRAINING", "false")
39+
t.Setenv("ENABLE_SQS_TERMINATION_DRAINING", "false")
40+
t.Setenv("ENABLE_REBALANCE_MONITORING", "true")
41+
t.Setenv("ENABLE_REBALANCE_DRAINING", "true")
42+
t.Setenv("GRACE_PERIOD", "12345")
43+
t.Setenv("IGNORE_DAEMON_SETS", "false")
44+
t.Setenv("KUBERNETES_SERVICE_HOST", "KUBERNETES_SERVICE_HOST")
45+
t.Setenv("KUBERNETES_SERVICE_PORT", "KUBERNETES_SERVICE_PORT")
46+
t.Setenv("NODE_NAME", "NODE_NAME")
47+
t.Setenv("NODE_TERMINATION_GRACE_PERIOD", "12345")
48+
t.Setenv("INSTANCE_METADATA_URL", "INSTANCE_METADATA_URL")
49+
t.Setenv("POD_TERMINATION_GRACE_PERIOD", "12345")
50+
t.Setenv("WEBHOOK_URL", "WEBHOOK_URL")
51+
t.Setenv("WEBHOOK_HEADERS", "WEBHOOK_HEADERS")
52+
t.Setenv("WEBHOOK_TEMPLATE", "WEBHOOK_TEMPLATE")
53+
t.Setenv("METADATA_TRIES", "100")
54+
t.Setenv("CORDON_ONLY", "false")
6655
nthConfig, err := config.ParseCliArgs()
6756
h.Ok(t, err)
6857

@@ -157,27 +146,27 @@ func TestParseCliArgsSuccess(t *testing.T) {
157146

158147
func TestParseCliArgsOverrides(t *testing.T) {
159148
resetFlagsForTest()
160-
setEnvForTest("USE_PROVIDER_ID", "true")
161-
setEnvForTest("DELETE_LOCAL_DATA", "true")
162-
setEnvForTest("DRY_RUN", "false")
163-
setEnvForTest("ENABLE_SCHEDULED_EVENT_DRAINING", "false")
164-
setEnvForTest("ENABLE_SPOT_INTERRUPTION_DRAINING", "true")
165-
setEnvForTest("ENABLE_SQS_TERMINATION_DRAINING", "false")
166-
setEnvForTest("ENABLE_REBALANCE_MONITORING", "true")
167-
setEnvForTest("ENABLE_REBALANCE_DRAINING", "true")
168-
setEnvForTest("GRACE_PERIOD", "99999")
169-
setEnvForTest("IGNORE_DAEMON_SETS", "true")
170-
setEnvForTest("KUBERNETES_SERVICE_HOST", "no")
171-
setEnvForTest("KUBERNETES_SERVICE_PORT", "no")
172-
setEnvForTest("NODE_NAME", "no")
173-
setEnvForTest("NODE_TERMINATION_GRACE_PERIOD", "99999")
174-
setEnvForTest("INSTANCE_METADATA_URL", "no")
175-
setEnvForTest("POD_TERMINATION_GRACE_PERIOD", "99999")
176-
setEnvForTest("WEBHOOK_URL", "no")
177-
setEnvForTest("WEBHOOK_HEADERS", "no")
178-
setEnvForTest("WEBHOOK_TEMPLATE", "no")
179-
setEnvForTest("METADATA_TRIES", "100")
180-
setEnvForTest("CORDON_ONLY", "true")
149+
t.Setenv("USE_PROVIDER_ID", "true")
150+
t.Setenv("DELETE_LOCAL_DATA", "true")
151+
t.Setenv("DRY_RUN", "false")
152+
t.Setenv("ENABLE_SCHEDULED_EVENT_DRAINING", "false")
153+
t.Setenv("ENABLE_SPOT_INTERRUPTION_DRAINING", "true")
154+
t.Setenv("ENABLE_SQS_TERMINATION_DRAINING", "false")
155+
t.Setenv("ENABLE_REBALANCE_MONITORING", "true")
156+
t.Setenv("ENABLE_REBALANCE_DRAINING", "true")
157+
t.Setenv("GRACE_PERIOD", "99999")
158+
t.Setenv("IGNORE_DAEMON_SETS", "true")
159+
t.Setenv("KUBERNETES_SERVICE_HOST", "no")
160+
t.Setenv("KUBERNETES_SERVICE_PORT", "no")
161+
t.Setenv("NODE_NAME", "no")
162+
t.Setenv("NODE_TERMINATION_GRACE_PERIOD", "99999")
163+
t.Setenv("INSTANCE_METADATA_URL", "no")
164+
t.Setenv("POD_TERMINATION_GRACE_PERIOD", "99999")
165+
t.Setenv("WEBHOOK_URL", "no")
166+
t.Setenv("WEBHOOK_HEADERS", "no")
167+
t.Setenv("WEBHOOK_TEMPLATE", "no")
168+
t.Setenv("METADATA_TRIES", "100")
169+
t.Setenv("CORDON_ONLY", "true")
181170
os.Args = []string{
182171
"cmd",
183172
"--use-provider-id=false",
@@ -238,9 +227,9 @@ func TestParseCliArgsOverrides(t *testing.T) {
238227

239228
func TestParseCliArgsWithGracePeriodSuccess(t *testing.T) {
240229
resetFlagsForTest()
241-
setEnvForTest("POD_TERMINATION_GRACE_PERIOD", "")
242-
setEnvForTest("NODE_NAME", "bla")
243-
setEnvForTest("GRACE_PERIOD", "12")
230+
t.Setenv("POD_TERMINATION_GRACE_PERIOD", "")
231+
t.Setenv("NODE_NAME", "bla")
232+
t.Setenv("GRACE_PERIOD", "12")
244233

245234
nthConfig, err := config.ParseCliArgs()
246235
h.Ok(t, err)
@@ -249,32 +238,32 @@ func TestParseCliArgsWithGracePeriodSuccess(t *testing.T) {
249238

250239
func TestParseCliArgsMissingNodeNameFailure(t *testing.T) {
251240
resetFlagsForTest()
252-
setEnvForTest("NODE_NAME", "")
241+
t.Setenv("NODE_NAME", "")
253242
_, err := config.ParseCliArgs()
254243
h.Assert(t, err != nil, "Failed to return error when node-name not provided")
255244
}
256245

257246
func TestParseCliArgsCreateFlagsFailure(t *testing.T) {
258247
resetFlagsForTest()
259-
setEnvForTest("DELETE_LOCAL_DATA", "something not true or false")
248+
t.Setenv("DELETE_LOCAL_DATA", "something not true or false")
260249
_, err := config.ParseCliArgs()
261250
h.Assert(t, err != nil, "Failed to return error when creating flags")
262251
}
263252

264253
func TestParseCliArgsAWSSession(t *testing.T) {
265254
resetFlagsForTest()
266-
setEnvForTest("ENABLE_SQS_TERMINATION_DRAINING", "true")
267-
setEnvForTest("AWS_REGION", "us-weast-1")
268-
setEnvForTest("NODE_NAME", "node")
255+
t.Setenv("ENABLE_SQS_TERMINATION_DRAINING", "true")
256+
t.Setenv("AWS_REGION", "us-weast-1")
257+
t.Setenv("NODE_NAME", "node")
269258
nthConfig, err := config.ParseCliArgs()
270259
h.Ok(t, err)
271260
h.Assert(t, nthConfig.AWSRegion == "us-weast-1", "Should find region as us-weast-1")
272261
}
273262

274263
func TestPrint_Human(t *testing.T) {
275264
resetFlagsForTest()
276-
setEnvForTest("NODE_NAME", "node")
277-
setEnvForTest("JSON_LOGGING", "false")
265+
t.Setenv("NODE_NAME", "node")
266+
t.Setenv("JSON_LOGGING", "false")
278267
nthConfig, err := config.ParseCliArgs()
279268
h.Ok(t, err)
280269
var printBuf bytes.Buffer
@@ -288,8 +277,8 @@ func TestPrint_Human(t *testing.T) {
288277

289278
func TestPrint_JSON(t *testing.T) {
290279
resetFlagsForTest()
291-
setEnvForTest("NODE_NAME", "node")
292-
setEnvForTest("JSON_LOGGING", "true")
280+
t.Setenv("NODE_NAME", "node")
281+
t.Setenv("JSON_LOGGING", "true")
293282
nthConfig, err := config.ParseCliArgs()
294283
h.Ok(t, err)
295284
var printBuf bytes.Buffer

0 commit comments

Comments
 (0)