github.com/tilt-dev/tilt@v0.33.15-0.20240515162809-0a22ed45d8a0/internal/testutils/env.go (about) 1 package testutils 2 3 import ( 4 "os" 5 "testing" 6 7 "github.com/stretchr/testify/require" 8 ) 9 10 // Setenv sets an environment variable for the duration of the test and restores the original value afterwards. 11 func Setenv(t testing.TB, key, value string) { 12 t.Helper() 13 captureAndRestoreEnv(t, key) 14 require.NoError(t, os.Setenv(key, value)) 15 } 16 17 // Unsetenv unsets an environment variable for the duration of the test and restores the original value afterwards. 18 func Unsetenv(t testing.TB, key string) { 19 t.Helper() 20 captureAndRestoreEnv(t, key) 21 require.NoError(t, os.Unsetenv(key)) 22 } 23 24 func captureAndRestoreEnv(t testing.TB, key string) { 25 t.Helper() 26 if curVal, ok := os.LookupEnv(key); ok { 27 t.Cleanup( 28 func() { 29 if err := os.Setenv(key, curVal); err != nil { 30 t.Errorf("Failed to restore env var %q: %v", key, err) 31 } 32 }, 33 ) 34 } else { 35 t.Cleanup( 36 func() { 37 if err := os.Unsetenv(key); err != nil { 38 t.Errorf("Failed to restore env var %q: %v", key, err) 39 } 40 }, 41 ) 42 } 43 }