Skip to content

Commit 80ad743

Browse files
committed
PatchEnvironment now correctly unsets env vars
os.GetEnv doesn't distinguish between unset vars and vars set to the empty string. However, LookupEnv does. This means PatchEnvironment would leave behind env vars set to the empty string when they were previously unset These env var would then be picked up by LookupEnv, which can cause unexpected errors.
1 parent 4ac48ac commit 80ad743

2 files changed

Lines changed: 37 additions & 7 deletions

File tree

patch.go

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -53,10 +53,14 @@ func PatchValue(dest, value interface{}) Restorer {
5353
// environment variable. A function is returned that will return the
5454
// environment to what it was before.
5555
func PatchEnvironment(name, value string) Restorer {
56-
oldValue := os.Getenv(name)
57-
os.Setenv(name, value)
56+
oldValue, oldValueSet := os.LookupEnv(name)
57+
_ = os.Setenv(name, value)
5858
return func() {
59-
os.Setenv(name, oldValue)
59+
if oldValueSet {
60+
_ = os.Setenv(name, oldValue)
61+
} else {
62+
_ = os.Unsetenv(name)
63+
}
6064
}
6165
}
6266

patch_test.go

Lines changed: 30 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -65,14 +65,40 @@ var _ = gc.Suite(&PatchEnvironmentSuite{})
6565
func (*PatchEnvironmentSuite) TestPatchEnvironment(c *gc.C) {
6666
const envName = "TESTING_PATCH_ENVIRONMENT"
6767
// remember the old value, and set it to something we can check
68-
oldValue := os.Getenv(envName)
69-
os.Setenv(envName, "initial")
68+
oldValue, oldValueSet := os.LookupEnv(envName)
69+
defer func() {
70+
if oldValueSet {
71+
_ = os.Setenv(envName, oldValue)
72+
} else {
73+
_ = os.Unsetenv(envName)
74+
}
75+
}()
76+
77+
_ = os.Setenv(envName, "initial")
7078
restore := testing.PatchEnvironment(envName, "new value")
7179
// Using check to make sure the environment gets set back properly in the test.
7280
c.Check(os.Getenv(envName), gc.Equals, "new value")
7381
restore()
7482
c.Check(os.Getenv(envName), gc.Equals, "initial")
75-
os.Setenv(envName, oldValue)
83+
}
84+
85+
func (*PatchEnvironmentSuite) TestPatchEnvironmentWithAbsentVar(c *gc.C) {
86+
const envName = "TESTING_PATCH_ENVIRONMENT"
87+
// remember the old value, and unset the var
88+
oldValue, oldValueSet := os.LookupEnv(envName)
89+
defer func() {
90+
if oldValueSet {
91+
_ = os.Setenv(envName, oldValue)
92+
}
93+
}()
94+
95+
_ = os.Unsetenv(envName)
96+
restore := testing.PatchEnvironment(envName, "new value")
97+
98+
c.Check(os.Getenv(envName), gc.Equals, "new value")
99+
restore()
100+
_, set := os.LookupEnv(envName)
101+
c.Check(set, gc.Equals, false)
76102
}
77103

78104
func (*PatchEnvironmentSuite) TestRestorerAdd(c *gc.C) {
@@ -89,7 +115,7 @@ func (*PatchEnvironmentSuite) TestPatchEnvPathPrepend(c *gc.C) {
89115
dir := "/bin/bar"
90116

91117
// just in case something goes wrong
92-
defer os.Setenv("PATH", oldPath)
118+
defer func() { _ = os.Setenv("PATH", oldPath) }()
93119

94120
restore := testing.PatchEnvPathPrepend(dir)
95121

0 commit comments

Comments
 (0)