Skip to content
This repository has been archived by the owner on Aug 9, 2021. It is now read-only.

Commit

Permalink
unset environmental varaibles, and restore them afterwards.
Browse files Browse the repository at this point in the history
  • Loading branch information
AnalogJ committed Sep 9, 2019
1 parent 4fea96a commit b6a78ea
Show file tree
Hide file tree
Showing 2 changed files with 70 additions and 0 deletions.
5 changes: 5 additions & 0 deletions pkg/config/config_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package config_test
import (
"github.com/analogj/capsulecd/pkg/config"
"github.com/analogj/capsulecd/pkg/pipeline"
"github.com/analogj/capsulecd/pkg/utils"
"github.com/stretchr/testify/require"
"os"
"path"
Expand All @@ -12,6 +13,10 @@ import (
func TestConfiguration_init_ShouldCorrectlyInitializeConfiguration(t *testing.T) {
t.Parallel()

//setup
defer utils.UnsetEnv("CAPSULE_")()


//test
testConfig, _ := config.Create()

Expand Down
65 changes: 65 additions & 0 deletions pkg/utils/env.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,65 @@
package utils

import (
"os"
"strings"
)

// UnsetEnv unsets all envars having prefix and returns a function
// that restores the env. Any newly added envars having prefix are
// also unset by restore. It is idiomatic to use with a defer.
//
// defer UnsetEnv("ACME_")()
//
// Note that modifying the env may have unpredictable results when
// tests are run with t.Parallel.
// NOTE: This is quick n' dirty from memory; write some tests for
// this code.
func UnsetEnv(prefix string) (restore func()) {
before := map[string]string{}

for _, e := range os.Environ() {
if !strings.HasPrefix(e, prefix) {
continue
}

parts := strings.SplitN(e, "=", 2)
before[parts[0]] = parts[1]

os.Unsetenv(parts[0])
}

return func() {
after := map[string]string{}

for _, e := range os.Environ() {
if !strings.HasPrefix(e, prefix) {
continue
}

parts := strings.SplitN(e, "=", 2)
after[parts[0]] = parts[1]

// Check if the envar previously existed
v, ok := before[parts[0]]
if !ok {
// This is a newly added envar with prefix, zap it
os.Unsetenv(parts[0])
continue
}

if parts[1] != v {
// If the envar value has changed, set it back
os.Setenv(parts[0], v)
}
}

// Still need to check if there have been any deleted envars
for k, v := range before {
if _, ok := after[k]; !ok {
// k is not present in after, so we set it.
os.Setenv(k, v)
}
}
}
}

0 comments on commit b6a78ea

Please sign in to comment.