github.com/aavshr/aws-sdk-go@v1.41.3/internal/sdktesting/env.go (about)

     1  package sdktesting
     2  
     3  import (
     4  	"os"
     5  	"runtime"
     6  	"strings"
     7  )
     8  
     9  // StashEnv stashes the current environment variables except variables listed in envToKeepx
    10  // Returns an function to pop out old environment
    11  func StashEnv(envToKeep ...string) func() {
    12  	if runtime.GOOS == "windows" {
    13  		envToKeep = append(envToKeep, "ComSpec")
    14  		envToKeep = append(envToKeep, "SYSTEM32")
    15  		envToKeep = append(envToKeep, "SYSTEMROOT")
    16  	}
    17  	envToKeep = append(envToKeep, "PATH")
    18  	extraEnv := getEnvs(envToKeep)
    19  	originalEnv := os.Environ()
    20  	os.Clearenv() // clear env
    21  	for key, val := range extraEnv {
    22  		os.Setenv(key, val)
    23  	}
    24  	return func() {
    25  		popEnv(originalEnv)
    26  	}
    27  }
    28  
    29  func getEnvs(envs []string) map[string]string {
    30  	extraEnvs := make(map[string]string)
    31  	for _, env := range envs {
    32  		if val, ok := os.LookupEnv(env); ok && len(val) > 0 {
    33  			extraEnvs[env] = val
    34  		}
    35  	}
    36  	return extraEnvs
    37  }
    38  
    39  // PopEnv takes the list of the environment values and injects them into the
    40  // process's environment variable data. Clears any existing environment values
    41  // that may already exist.
    42  func popEnv(env []string) {
    43  	os.Clearenv()
    44  
    45  	for _, e := range env {
    46  		p := strings.SplitN(e, "=", 2)
    47  		k, v := p[0], ""
    48  		if len(p) > 1 {
    49  			v = p[1]
    50  		}
    51  		os.Setenv(k, v)
    52  	}
    53  }