env.go

Documentation: gotest.tools/v3/env

     1  /*
     2  Package env provides functions to test code that read environment variables
     3  or the current working directory.
     4  */
     5  package env // import "gotest.tools/v3/env"
     6  
     7  import (
     8  	"os"
     9  	"strings"
    10  
    11  	"gotest.tools/v3/assert"
    12  	"gotest.tools/v3/internal/cleanup"
    13  )
    14  
    15  type helperT interface {
    16  	Helper()
    17  }
    18  
    19  // Patch changes the value of an environment variable, and returns a
    20  // function which will reset the the value of that variable back to the
    21  // previous state.
    22  //
    23  // When used with Go 1.14+ the unpatch function will be called automatically
    24  // when the test ends, unless the TEST_NOCLEANUP env var is set to true.
    25  //
    26  // Deprecated: use t.SetEnv
    27  func Patch(t assert.TestingT, key, value string) func() {
    28  	if ht, ok := t.(helperT); ok {
    29  		ht.Helper()
    30  	}
    31  	oldValue, envVarExists := os.LookupEnv(key)
    32  	assert.NilError(t, os.Setenv(key, value))
    33  	clean := func() {
    34  		if ht, ok := t.(helperT); ok {
    35  			ht.Helper()
    36  		}
    37  		if !envVarExists {
    38  			assert.NilError(t, os.Unsetenv(key))
    39  			return
    40  		}
    41  		assert.NilError(t, os.Setenv(key, oldValue))
    42  	}
    43  	cleanup.Cleanup(t, clean)
    44  	return clean
    45  }
    46  
    47  // PatchAll sets the environment to env, and returns a function which will
    48  // reset the environment back to the previous state.
    49  //
    50  // When used with Go 1.14+ the unpatch function will be called automatically
    51  // when the test ends, unless the TEST_NOCLEANUP env var is set to true.
    52  func PatchAll(t assert.TestingT, env map[string]string) func() {
    53  	if ht, ok := t.(helperT); ok {
    54  		ht.Helper()
    55  	}
    56  	oldEnv := os.Environ()
    57  	os.Clearenv()
    58  
    59  	for key, value := range env {
    60  		assert.NilError(t, os.Setenv(key, value), "setenv %s=%s", key, value)
    61  	}
    62  	clean := func() {
    63  		if ht, ok := t.(helperT); ok {
    64  			ht.Helper()
    65  		}
    66  		os.Clearenv()
    67  		for key, oldVal := range ToMap(oldEnv) {
    68  			assert.NilError(t, os.Setenv(key, oldVal), "setenv %s=%s", key, oldVal)
    69  		}
    70  	}
    71  	cleanup.Cleanup(t, clean)
    72  	return clean
    73  }
    74  
    75  // ToMap takes a list of strings in the format returned by [os.Environ] and
    76  // returns a mapping of keys to values.
    77  func ToMap(env []string) map[string]string {
    78  	result := map[string]string{}
    79  	for _, raw := range env {
    80  		key, value := getParts(raw)
    81  		result[key] = value
    82  	}
    83  	return result
    84  }
    85  
    86  func getParts(raw string) (string, string) {
    87  	if raw == "" {
    88  		return "", ""
    89  	}
    90  	// Environment variables on windows can begin with =
    91  	// http://blogs.msdn.com/b/oldnewthing/archive/2010/05/06/10008132.aspx
    92  	parts := strings.SplitN(raw[1:], "=", 2)
    93  	key := raw[:1] + parts[0]
    94  	if len(parts) == 1 {
    95  		return key, ""
    96  	}
    97  	return key, parts[1]
    98  }
    99  
   100  // ChangeWorkingDir to the directory, and return a function which restores the
   101  // previous working directory.
   102  //
   103  // When used with Go 1.14+ the previous working directory will be restored
   104  // automatically when the test ends, unless the TEST_NOCLEANUP env var is set to
   105  // true.
   106  func ChangeWorkingDir(t assert.TestingT, dir string) func() {
   107  	if ht, ok := t.(helperT); ok {
   108  		ht.Helper()
   109  	}
   110  	cwd, err := os.Getwd()
   111  	assert.NilError(t, err)
   112  	assert.NilError(t, os.Chdir(dir))
   113  	clean := func() {
   114  		if ht, ok := t.(helperT); ok {
   115  			ht.Helper()
   116  		}
   117  		assert.NilError(t, os.Chdir(cwd))
   118  	}
   119  	cleanup.Cleanup(t, clean)
   120  	return clean
   121  }
   122
View as plain text