Refactor shell and add function to return output struct (#51)

* Refactor shell and add function to return output struct

* Split out the stdlib monkeypatches to a different file

* Add stream version of returning output struct

* Port terratest version of output streaming for more reliability

* Make sure newlines are included properly

* Add sleep to ensure stdout is read before stderr

* We need to Println in streammode

* Add comments about the relation to terratest

* Remove unused imports

Author: yorinasub17

shell/cmd.go

@@ -1,9 +1,7 @@
 package shell
 
 import (
-	"bufio"
 	"fmt"
-	"io"
 	"os"
 	"os/exec"
 	"strings"
@@ -27,74 +25,47 @@ func RunShellCommand(options *ShellOptions, command string, args ...string) erro
 	return errors.WithStackTrace(cmd.Run())
 }
 
+// Run the specified shell command with the specified arguments. Return its stdout, stderr, and interleaved output as
+// separate strings in a struct.
+func RunShellCommandAndGetOutputStruct(options *ShellOptions, command string, args ...string) (*Output, error) {
+	return runShellCommand(options, false, command, args...)
+}
+
 // Run the specified shell command with the specified arguments. Return its stdout and stderr as a string
 func RunShellCommandAndGetOutput(options *ShellOptions, command string, args ...string) (string, error) {
-	logCommand(options, command, args...)
-	cmd := exec.Command(command, args...)
-
-	cmd.Stdin = os.Stdin
-
-	setCommandOptions(options, cmd)
-
-	out, err := cmd.CombinedOutput()
-	return string(out), errors.WithStackTrace(err)
+	out, err := runShellCommand(options, false, command, args...)
+	return out.Combined(), err
 }
 
-// Run the specified shell command with the specified arguments. Return its stdout and stderr as a string and also
-// stream stdout and stderr to the OS stdout/stderr
+// Run the specified shell command with the specified arguments. Return its interleaved stdout and stderr as a string
+// and also stream stdout and stderr to the OS stdout/stderr
 func RunShellCommandAndGetAndStreamOutput(options *ShellOptions, command string, args ...string) (string, error) {
-	logCommand(options, command, args...)
-	cmd := exec.Command(command, args...)
-
-	setCommandOptions(options, cmd)
-
-	cmd.Stdin = os.Stdin
-
-	stdout, err := cmd.StdoutPipe()
-	if err != nil {
-		return "", errors.WithStackTrace(err)
-	}
-
-	stderr, err := cmd.StderrPipe()
-	if err != nil {
-		return "", errors.WithStackTrace(err)
-	}
-
-	if err := cmd.Start(); err != nil {
-		return "", errors.WithStackTrace(err)
-	}
-
-	output, err := readStdoutAndStderr(
-		stdout,
-		true,
-		stderr,
-		true,
-		options,
-	)
-	if err != nil {
-		return output, err
-	}
-
-	err = cmd.Wait()
-	return output, errors.WithStackTrace(err)
+	out, err := runShellCommand(options, true, command, args...)
+	return out.Combined(), err
 }
 
 // Run the specified shell command with the specified arguments. Return its stdout as a string
 func RunShellCommandAndGetStdout(options *ShellOptions, command string, args ...string) (string, error) {
-	logCommand(options, command, args...)
-	cmd := exec.Command(command, args...)
-
-	cmd.Stdin = os.Stdin
-
-	setCommandOptions(options, cmd)
-
-	out, err := cmd.Output()
-	return string(out), errors.WithStackTrace(err)
+	out, err := runShellCommand(options, false, command, args...)
+	return out.Stdout(), err
 }
 
 // Run the specified shell command with the specified arguments. Return its stdout as a string and also stream stdout
 // and stderr to the OS stdout/stderr
 func RunShellCommandAndGetStdoutAndStreamOutput(options *ShellOptions, command string, args ...string) (string, error) {
+	out, err := runShellCommand(options, true, command, args...)
+	return out.Stdout(), err
+}
+
+// Run the specified shell command with the specified arguments. Return its stdout, stderr, and interleaved output as a
+// struct and also stream stdout and stderr to the OS stdout/stderr
+func RunShellCommandAndGetOutputStructAndStreamOutput(options *ShellOptions, command string, args ...string) (*Output, error) {
+	return runShellCommand(options, true, command, args...)
+}
+
+// Run the specified shell command with the specified arguments. Return its stdout and stderr as a string and also
+// stream stdout and stderr to the OS stdout/stderr
+func runShellCommand(options *ShellOptions, streamOutput bool, command string, args ...string) (*Output, error) {
 	logCommand(options, command, args...)
 	cmd := exec.Command(command, args...)
 
@@ -104,24 +75,23 @@ func RunShellCommandAndGetStdoutAndStreamOutput(options *ShellOptions, command s
 
 	stdout, err := cmd.StdoutPipe()
 	if err != nil {
-		return "", errors.WithStackTrace(err)
+		return nil, errors.WithStackTrace(err)
 	}
 
 	stderr, err := cmd.StderrPipe()
 	if err != nil {
-		return "", errors.WithStackTrace(err)
+		return nil, errors.WithStackTrace(err)
 	}
 
 	if err := cmd.Start(); err != nil {
-		return "", errors.WithStackTrace(err)
+		return nil, errors.WithStackTrace(err)
 	}
 
 	output, err := readStdoutAndStderr(
+		options.Logger,
+		streamOutput,
 		stdout,
-		true,
 		stderr,
-		false,
-		options,
 	)
 	if err != nil {
 		return output, err
@@ -131,48 +101,6 @@ func RunShellCommandAndGetStdoutAndStreamOutput(options *ShellOptions, command s
 	return output, errors.WithStackTrace(err)
 }
 
-// This function captures stdout and stderr while still printing it to the stdout and stderr of this Go program
-func readStdoutAndStderr(
-	stdout io.ReadCloser,
-	includeStdout bool,
-	stderr io.ReadCloser,
-	includeStderr bool,
-	options *ShellOptions,
-) (string, error) {
-	allOutput := []string{}
-
-	stdoutScanner := bufio.NewScanner(stdout)
-	stderrScanner := bufio.NewScanner(stderr)
-
-	for {
-		if stdoutScanner.Scan() {
-			text := stdoutScanner.Text()
-			options.Logger.Println(text)
-			if includeStdout {
-				allOutput = append(allOutput, text)
-			}
-		} else if stderrScanner.Scan() {
-			text := stderrScanner.Text()
-			options.Logger.Println(text)
-			if includeStderr {
-				allOutput = append(allOutput, text)
-			}
-		} else {
-			break
-		}
-	}
-
-	if err := stdoutScanner.Err(); err != nil {
-		return "", errors.WithStackTrace(err)
-	}
-
-	if err := stderrScanner.Err(); err != nil {
-		return "", errors.WithStackTrace(err)
-	}
-
-	return strings.Join(allOutput, "\n"), nil
-}
-
 func logCommand(options *ShellOptions, command string, args ...string) {
 	if options.SensitiveArgs {
 		options.Logger.Infof("Running command: %s (args redacted)", command)

shell/cmd_test.go

@@ -30,6 +30,14 @@ func TestRunShellCommandAndGetOutput(t *testing.T) {
 	assert.Equal(t, "hi\n", out)
 }
 
+func TestRunShellCommandAndGetOutputNoTerminatingNewLine(t *testing.T) {
+	t.Parallel()
+
+	out, err := RunShellCommandAndGetOutput(NewShellOptions(), "echo", "-n", "hi")
+	assert.NoError(t, err)
+	assert.Equal(t, "hi", out)
+}
+
 func TestRunShellCommandAndGetStdoutReturnsStdout(t *testing.T) {
 	t.Parallel()
 
@@ -46,6 +54,32 @@ func TestRunShellCommandAndGetStdoutDoesNotReturnStderr(t *testing.T) {
 	assert.Equal(t, "", out)
 }
 
+func TestRunShellCommandAndGetStdoutAndStreamOutputDoesNotReturnStderr(t *testing.T) {
+	t.Parallel()
+
+	out, err := RunShellCommandAndGetStdoutAndStreamOutput(NewShellOptions(), filepath.Join("test-fixture", "echo_hi_stderr.sh"))
+	assert.NoError(t, err)
+	assert.Equal(t, "", out)
+}
+
+func TestRunShellCommandAndGetAndStreamOutput(t *testing.T) {
+	t.Parallel()
+
+	out, err := RunShellCommandAndGetAndStreamOutput(NewShellOptions(), filepath.Join("test-fixture", "echo_stdoutstderr.sh"))
+	assert.NoError(t, err)
+	assert.Equal(t, "hello\nworld\n", out)
+}
+
+func TestRunShellCommandAndGetOutputStruct(t *testing.T) {
+	t.Parallel()
+
+	out, err := RunShellCommandAndGetOutputStruct(NewShellOptions(), filepath.Join("test-fixture", "echo_stdoutstderr.sh"))
+	assert.NoError(t, err)
+	assert.Equal(t, "hello\nworld\n", out.Combined())
+	assert.Equal(t, "hello\n", out.Stdout())
+	assert.Equal(t, "world\n", out.Stderr())
+}
+
 func TestRunShellCommandWithEnv(t *testing.T) {
 	t.Parallel()