Fix running a command with a new env variable

Previously, if the variable was not already existing, the variable was dropped.

Author: guillaumerose

pkg/os/exec.go

@@ -14,7 +14,7 @@ func runCmd(command string, args []string, env map[string]string) (string, strin
 	if len(env) != 0 {
 		cmd.Env = os.Environ()
 		for key, value := range env {
-			cmd.Env = ReplaceEnv(cmd.Env, key, value)
+			cmd.Env = ReplaceOrAddEnv(cmd.Env, key, value)
 		}
 	}
 	stdOut := new(bytes.Buffer)
@@ -52,9 +52,25 @@ func RunWithPrivilege(reason string, cmdAndArgs ...string) (string, string, erro
 }
 
 func RunWithDefaultLocale(command string, args ...string) (string, string, error) {
-	return run(command, args, map[string]string{"LC_ALL": "C", "LANG": "C"})
+	return RunWithDefaultLocaleAndEnv(command, args, map[string]string{})
 }
 
 func RunWithDefaultLocalePrivate(command string, args ...string) (string, string, error) {
-	return runPrivate(command, args, map[string]string{"LC_ALL": "C", "LANG": "C"})
+	return RunWithDefaultLocalePrivateAndEnv(command, args, map[string]string{})
+}
+
+func RunWithDefaultLocaleAndEnv(command string, args []string, extraEnv map[string]string) (string, string, error) {
+	env := map[string]string{"LC_ALL": "C", "LANG": "C"}
+	for k, v := range extraEnv {
+		env[k] = v
+	}
+	return run(command, args, env)
+}
+
+func RunWithDefaultLocalePrivateAndEnv(command string, args []string, extraEnv map[string]string) (string, string, error) {
+	env := map[string]string{"LC_ALL": "C", "LANG": "C"}
+	for k, v := range extraEnv {
+		env[k] = v
+	}
+	return runPrivate(command, args, env)
 }

pkg/os/util.go

@@ -37,19 +37,25 @@ func CurrentOS() OS {
 	panic("Unexpected OS type")
 }
 
-// ReplaceEnv changes the value of an environment variable
+// ReplaceOrAddEnv changes the value of an environment variable if it exists otherwise add the new variable
 // It drops the existing value and appends the new value in-place
-func ReplaceEnv(variables []string, varName string, value string) []string {
+func ReplaceOrAddEnv(variables []string, varName string, value string) []string {
 	var result []string
+
+	found := false
 	for _, e := range variables {
 		pair := strings.Split(e, "=")
 		if pair[0] != varName {
 			result = append(result, e)
 		} else {
+			found = true
 			result = append(result, fmt.Sprintf("%s=%s", varName, value))
 		}
 	}
 
+	if !found {
+		result = append(result, fmt.Sprintf("%s=%s", varName, value))
+	}
 	return result
 }
 

pkg/os/util_test.go

@@ -0,0 +1,24 @@
+package os
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestReplaceEnv(t *testing.T) {
+	env := []string{"HOME=/home/user/", "PATH=/bin:/sbin:/usr/bin", "LC_ALL=de_DE.UTF8"}
+	replaced := ReplaceOrAddEnv(env, "LC_ALL", "C")
+
+	assert.Len(t, replaced, 3)
+	assert.Equal(t, env[2], "LC_ALL=de_DE.UTF8")
+	assert.Equal(t, replaced[2], "LC_ALL=C")
+}
+
+func TestAddEnv(t *testing.T) {
+	env := []string{"HOME=/home/user/", "PATH=/bin:/sbin:/usr/bin", "LC_ALL=de_DE.UTF8"}
+	replaced := ReplaceOrAddEnv(env, "KUBECONFIG", "some-data")
+
+	assert.Len(t, replaced, 4)
+	assert.Equal(t, replaced[3], "KUBECONFIG=some-data")
+}