Support custom environment variables when running commands via the shell module

shell/cmd.go

@@ -27,7 +27,7 @@ func RunShellCommand(options *ShellOptions, command string, args ...string) erro
 	cmd.Stdout = os.Stdout
 	cmd.Stderr = os.Stderr
 
-	cmd.Dir = options.WorkingDir
+	setCommandOptions(options, cmd)
 
 	return errors.WithStackTrace(cmd.Run())
 }
@@ -43,7 +43,8 @@ func RunShellCommandAndGetOutput(options *ShellOptions, command string, args ...
 	cmd := exec.Command(command, args...)
 
 	cmd.Stdin = os.Stdin
-	cmd.Dir = options.WorkingDir
+
+	setCommandOptions(options, cmd)
 
 	out, err := cmd.CombinedOutput()
 	return string(out), errors.WithStackTrace(err)
@@ -60,7 +61,7 @@ func RunShellCommandAndGetAndStreamOutput(options *ShellOptions, command string,
 
 	cmd := exec.Command(command, args...)
 
-	cmd.Dir = options.WorkingDir
+	setCommandOptions(options, cmd)
 
 	cmd.Stdin = os.Stdin
 
@@ -133,3 +134,20 @@ func CommandInstalledE(command string) error {
 	}
 	return nil
 }
+
+// setCommandOptions takes the shell options and maps them to the configurations for the exec.Cmd object, applying them
+// to the passed in Cmd object.
+func setCommandOptions(options *ShellOptions, cmd *exec.Cmd) {
+	cmd.Dir = options.WorkingDir
+	cmd.Env = formatEnvVars(options)
+}
+
+// formatEnvVars takes environment variables encoded into ShellOptions and converts them to a format understood by
+// exec.Command
+func formatEnvVars(options *ShellOptions) []string {
+	env := os.Environ()
+	for key, value := range options.Env {
+		env = append(env, fmt.Sprintf("%s=%s", key, value))
+	}
+	return env
+}

shell/cmd_test.go

@@ -2,32 +2,52 @@ package shell
 
 import (
 	"bytes"
-	"github.com/stretchr/testify/assert"
+	"fmt"
 	"testing"
 
 	"github.com/gruntwork-io/gruntwork-cli/logging"
+	"github.com/stretchr/testify/assert"
 )
 
 func TestRunShellCommand(t *testing.T) {
 	t.Parallel()
 
-	assert.Nil(t, RunShellCommand(NewShellOptions(), "echo", "hi"))
+	assert.NoError(t, RunShellCommand(NewShellOptions(), "echo", "hi"))
 }
 
 func TestRunShellCommandInvalidCommand(t *testing.T) {
 	t.Parallel()
 
-	assert.NotNil(t, RunShellCommand(NewShellOptions(), "not-a-real-command"))
+	assert.Error(t, RunShellCommand(NewShellOptions(), "not-a-real-command"))
 }
 
 func TestRunShellCommandAndGetOutput(t *testing.T) {
 	t.Parallel()
 
 	out, err := RunShellCommandAndGetOutput(NewShellOptions(), "echo", "hi")
-	assert.Nil(t, err, "Unexpected error: %v", err)
+	assert.NoError(t, err)
 	assert.Equal(t, "hi\n", out)
 }
 
+func TestRunShellCommandWithEnv(t *testing.T) {
+	t.Parallel()
+
+	envVars := map[string]string{
+		"TEST_WITH_SPACES":  "test with spaces",
+		"TEST_WITH_EQUALS":  "test=with=equals",
+		"TEST_START_EQUALS": "=teststartequals",
+		"TEST_BLANK":        "",
+	}
+	options := NewShellOptions()
+	options.Env = envVars
+
+	for k, v := range envVars {
+		out, err := RunShellCommandAndGetOutput(options, "bash", "-c", fmt.Sprintf("echo $%s", k))
+		assert.NoError(t, err)
+		assert.Equal(t, fmt.Sprintf("%s\n", v), out)
+	}
+}
+
 func TestCommandInstalledOnValidCommand(t *testing.T) {
 	t.Parallel()
 
@@ -43,13 +63,13 @@ func TestCommandInstalledOnInvalidCommand(t *testing.T) {
 func TestCommandInstalledEOnValidCommand(t *testing.T) {
 	t.Parallel()
 
-	assert.Nil(t, CommandInstalledE("echo"))
+	assert.NoError(t, CommandInstalledE("echo"))
 }
 
 func TestCommandInstalledEOnInvalidCommand(t *testing.T) {
 	t.Parallel()
 
-	assert.NotNil(t, CommandInstalledE("not-a-real-command"))
+	assert.Error(t, CommandInstalledE("not-a-real-command"))
 }
 
 // Test that when SensitiveArgs is true, do not log the args
@@ -63,7 +83,7 @@ func TestSensitiveArgsTrueHidesOnRunShellCommand(t *testing.T) {
 	options.SensitiveArgs = true
 	options.Logger = logger
 
-	assert.Nil(t, RunShellCommand(options, "echo", "hi"))
+	assert.NoError(t, RunShellCommand(options, "echo", "hi"))
 	assert.NotContains(t, buffer.String(), "hi")
 	assert.Contains(t, buffer.String(), "echo")
 }
@@ -78,7 +98,7 @@ func TestSensitiveArgsFalseShowsOnRunShellCommand(t *testing.T) {
 	options := NewShellOptions()
 	options.Logger = logger
 
-	assert.Nil(t, RunShellCommand(options, "echo", "hi"))
+	assert.NoError(t, RunShellCommand(options, "echo", "hi"))
 	assert.Contains(t, buffer.String(), "hi")
 	assert.Contains(t, buffer.String(), "echo")
 }
@@ -95,7 +115,7 @@ func TestSensitiveArgsTrueHidesOnRunShellCommandAndGetOutput(t *testing.T) {
 	options.Logger = logger
 
 	_, err := RunShellCommandAndGetOutput(options, "echo", "hi")
-	assert.Nil(t, err)
+	assert.NoError(t, err)
 	assert.NotContains(t, buffer.String(), "hi")
 	assert.Contains(t, buffer.String(), "echo")
 }
@@ -111,7 +131,7 @@ func TestSensitiveArgsFalseShowsOnRunShellCommandAndGetOutput(t *testing.T) {
 	options.Logger = logger
 
 	_, err := RunShellCommandAndGetOutput(options, "echo", "hi")
-	assert.Nil(t, err)
+	assert.NoError(t, err)
 	assert.Contains(t, buffer.String(), "hi")
 	assert.Contains(t, buffer.String(), "echo")
 }

shell/options.go

@@ -9,7 +9,8 @@ type ShellOptions struct {
 	NonInteractive bool
 	Logger         *logrus.Logger
 	WorkingDir     string
-	SensitiveArgs  bool // If true, will not log the arguments to the command
+	SensitiveArgs  bool              // If true, will not log the arguments to the command
+	Env            map[string]string // Additional environment variables to set
 }
 
 func NewShellOptions() *ShellOptions {
@@ -18,5 +19,6 @@ func NewShellOptions() *ShellOptions {
 		Logger:         logging.GetLogger(""),
 		WorkingDir:     ".",
 		SensitiveArgs:  false,
+		Env:            map[string]string{},
 	}
 }