|
| 1 | +package utils |
| 2 | + |
| 3 | +import ( |
| 4 | + "os" |
| 5 | + "strings" |
| 6 | +) |
| 7 | + |
| 8 | +// UnsetEnv unsets all envars having prefix and returns a function |
| 9 | +// that restores the env. Any newly added envars having prefix are |
| 10 | +// also unset by restore. It is idiomatic to use with a defer. |
| 11 | +// |
| 12 | +// defer UnsetEnv("ACME_")() |
| 13 | +// |
| 14 | +// Note that modifying the env may have unpredictable results when |
| 15 | +// tests are run with t.Parallel. |
| 16 | +// NOTE: This is quick n' dirty from memory; write some tests for |
| 17 | +// this code. |
| 18 | +func UnsetEnv(prefix string) (restore func()) { |
| 19 | + before := map[string]string{} |
| 20 | + |
| 21 | + for _, e := range os.Environ() { |
| 22 | + if !strings.HasPrefix(e, prefix) { |
| 23 | + continue |
| 24 | + } |
| 25 | + |
| 26 | + parts := strings.SplitN(e, "=", 2) |
| 27 | + before[parts[0]] = parts[1] |
| 28 | + |
| 29 | + os.Unsetenv(parts[0]) |
| 30 | + } |
| 31 | + |
| 32 | + return func() { |
| 33 | + after := map[string]string{} |
| 34 | + |
| 35 | + for _, e := range os.Environ() { |
| 36 | + if !strings.HasPrefix(e, prefix) { |
| 37 | + continue |
| 38 | + } |
| 39 | + |
| 40 | + parts := strings.SplitN(e, "=", 2) |
| 41 | + after[parts[0]] = parts[1] |
| 42 | + |
| 43 | + // Check if the envar previously existed |
| 44 | + v, ok := before[parts[0]] |
| 45 | + if !ok { |
| 46 | + // This is a newly added envar with prefix, zap it |
| 47 | + os.Unsetenv(parts[0]) |
| 48 | + continue |
| 49 | + } |
| 50 | + |
| 51 | + if parts[1] != v { |
| 52 | + // If the envar value has changed, set it back |
| 53 | + os.Setenv(parts[0], v) |
| 54 | + } |
| 55 | + } |
| 56 | + |
| 57 | + // Still need to check if there have been any deleted envars |
| 58 | + for k, v := range before { |
| 59 | + if _, ok := after[k]; !ok { |
| 60 | + // k is not present in after, so we set it. |
| 61 | + os.Setenv(k, v) |
| 62 | + } |
| 63 | + } |
| 64 | + } |
| 65 | +} |
0 commit comments