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  }