Merge pull request #9 from SimonBaeumer/add-multiplexed-writer

Add multiplexed writer

README.md

@@ -33,9 +33,13 @@ To configure the command a option function will be passed which receives the com
 Default option functions:
 
  - `cmd.WithStandardStreams`
+ - `cmd.WithCustomStdout(...io.Writers)`
+ - `cmd.WithCustomStderr(...io.Writers)`
  - `cmd.WithTimeout(time.Duration)`
  - `cmd.WithoutTimeout`
  - `cmd.WithWorkingDir(string)`
+ - `cmd.WithEnvironmentVariables(cmd.EnvVars)`
+ - `cmd.WithInheritedEnvironment(cmd.EnvVars)`
 
 #### Example
 
@@ -55,6 +59,23 @@ c := cmd.NewCommand("pwd", setWorkingDir)
 c.Execute()
 ```
 
+### Testing
+
+You can catch output streams to `stdout` and `stderr` with `cmd.CaptureStandardOut`. 
+
+```golang
+// caputred is the captured output from all executed source code
+// fnResult contains the result of the executed function
+captured, fnResult := cmd.CaptureStandardOut(func() interface{} {
+    c := NewCommand("echo hello", cmd.WithStandardStream)
+    err := c.Execute()
+    return err
+})
+
+// prints "hello"
+fmt.Println(captured)
+```
+
 ## Development
 
 ### Running tests

command.go

@@ -10,7 +10,7 @@ import (
 	"time"
 )
 
-//Command represents a single command which can be executed
+// Command represents a single command which can be executed
 type Command struct {
 	Command      string
 	Env          []string
@@ -22,10 +22,20 @@ type Command struct {
 	executed     bool
 	exitCode     int
 	// stderr and stdout retrieve the output after the command was executed
-	stderr bytes.Buffer
-	stdout bytes.Buffer
+	stderr   bytes.Buffer
+	stdout   bytes.Buffer
+	combined bytes.Buffer
 }
 
+// EnvVars represents a map where the key is the name of the env variable
+// and the value is the value of the variable
+//
+// Example:
+//
+//  env := map[string]string{"ENV": "VALUE"}
+//
+type EnvVars map[string]string
+
 // NewCommand creates a new command
 // You can add option with variadic option argument
 // Default timeout is set to 30 minutes
@@ -68,8 +78,28 @@ func NewCommand(cmd string, options ...func(*Command)) *Command {
 //     c.Execute()
 //
 func WithStandardStreams(c *Command) {
-	c.StdoutWriter = os.Stdout
-	c.StderrWriter = os.Stderr
+	c.StdoutWriter = NewMultiplexedWriter(os.Stdout, &c.stdout, &c.combined)
+	c.StderrWriter = NewMultiplexedWriter(os.Stderr, &c.stdout, &c.combined)
+}
+
+// WithCustomStdout allows to add custom writers to stdout
+func WithCustomStdout(writers ...io.Writer) func(c *Command) {
+	return func(c *Command) {
+		writers = append(writers, &c.stdout, &c.combined)
+		c.StdoutWriter = NewMultiplexedWriter(writers...)
+
+		c.StderrWriter = NewMultiplexedWriter(&c.stderr, &c.combined)
+	}
+}
+
+// WithCustomStderr allows to add custom writers to stderr
+func WithCustomStderr(writers ...io.Writer) func(c *Command) {
+	return func(c *Command) {
+		writers = append(writers, &c.stderr, &c.combined)
+		c.StderrWriter = NewMultiplexedWriter(writers...)
+
+		c.StdoutWriter = NewMultiplexedWriter(&c.stdout, &c.combined)
+	}
 }
 
 // WithTimeout sets the timeout of the command
@@ -95,25 +125,52 @@ func WithWorkingDir(dir string) func(c *Command) {
 	}
 }
 
+// WithInheritedEnvironment uses the env from the current process and
+// allow to add more variables.
+func WithInheritedEnvironment(env EnvVars) func(c *Command) {
+	return func(c *Command) {
+		c.Env = os.Environ()
+
+		// Set custom variables
+		fn := WithEnvironmentVariables(env)
+		fn(c)
+	}
+}
+
+// WithEnvironmentVariables sets environment variables for the executed command
+func WithEnvironmentVariables(env EnvVars) func(c *Command) {
+	return func(c *Command) {
+		for key, value := range env {
+			c.AddEnv(key, value)
+		}
+	}
+}
+
 // AddEnv adds an environment variable to the command
 // If a variable gets passed like ${VAR_NAME} the env variable will be read out by the current shell
 func (c *Command) AddEnv(key string, value string) {
 	value = os.ExpandEnv(value)
 	c.Env = append(c.Env, fmt.Sprintf("%s=%s", key, value))
 }
 
-//Stdout returns the output to stdout
+// Stdout returns the output to stdout
 func (c *Command) Stdout() string {
 	c.isExecuted("Stdout")
 	return c.stdout.String()
 }
 
-//Stderr returns the output to stderr
+// Stderr returns the output to stderr
 func (c *Command) Stderr() string {
 	c.isExecuted("Stderr")
 	return c.stderr.String()
 }
 
+// Combined returns the combined output of stderr and stdout according to their timeline
+func (c *Command) Combined() string {
+	c.isExecuted("Combined")
+	return c.combined.String()
+}
+
 //ExitCode returns the exit code of the command
 func (c *Command) ExitCode() int {
 	c.isExecuted("ExitCode")

command_linux_test.go

@@ -1,6 +1,7 @@
 package cmd
 
 import (
+	"bytes"
 	"github.com/stretchr/testify/assert"
 	"io/ioutil"
 	"os"
@@ -81,3 +82,59 @@ func TestCommand_WithInvalidDir(t *testing.T) {
 	assert.NotNil(t, err)
 	assert.Equal(t, "chdir /invalid: no such file or directory", err.Error())
 }
+
+func TestWithInheritedEnvironment(t *testing.T) {
+	os.Setenv("FROM_OS", "is on os")
+	os.Setenv("OVERWRITE", "is on os but should be overwritten")
+	defer func() {
+		os.Unsetenv("FROM_OS")
+		os.Unsetenv("OVERWRITE")
+	}()
+
+	c := NewCommand(
+		"echo $FROM_OS $OVERWRITE",
+		WithInheritedEnvironment(map[string]string{"OVERWRITE": "overwritten"}))
+	c.Execute()
+
+	assertEqualWithLineBreak(t, "is on os overwritten", c.Stdout())
+}
+
+func TestWithCustomStderr(t *testing.T) {
+	writer := bytes.Buffer{}
+	c := NewCommand(">&2 echo stderr; sleep 0.01; echo stdout;", WithCustomStderr(&writer))
+	c.Execute()
+
+	assertEqualWithLineBreak(t, "stderr", writer.String())
+	assertEqualWithLineBreak(t, "stdout", c.Stdout())
+	assertEqualWithLineBreak(t, "stderr", c.Stderr())
+	assertEqualWithLineBreak(t, "stderr\nstdout", c.Combined())
+}
+
+func TestWithCustomStdout(t *testing.T) {
+	writer := bytes.Buffer{}
+	c := NewCommand(">&2 echo stderr; sleep 0.01; echo stdout;", WithCustomStdout(&writer))
+	c.Execute()
+
+	assertEqualWithLineBreak(t, "stdout", writer.String())
+	assertEqualWithLineBreak(t, "stdout", c.Stdout())
+	assertEqualWithLineBreak(t, "stderr", c.Stderr())
+	assertEqualWithLineBreak(t, "stderr\nstdout", c.Combined())
+}
+
+func TestWithStandardStreams(t *testing.T) {
+	out, err := CaptureStandardOutput(func() interface{} {
+		c := NewCommand(">&2 echo stderr; sleep 0.01; echo stdout;", WithStandardStreams)
+		err := c.Execute()
+		return err
+	})
+
+	assertEqualWithLineBreak(t, "stderr\nstdout", out)
+	assert.Nil(t, err)
+}
+
+func TestWithEnvironmentVariables(t *testing.T) {
+	c := NewCommand("echo $env", WithEnvironmentVariables(map[string]string{"env": "value"}))
+	c.Execute()
+
+	assertEqualWithLineBreak(t, "value", c.Stdout())
+}

command_test.go

@@ -126,13 +126,3 @@ func TestCommand_SetOptions(t *testing.T) {
 	assert.Equal(t, time.Duration(1000000000), c.Timeout)
 	assertEqualWithLineBreak(t, "test", writer.String())
 }
-
-func assertEqualWithLineBreak(t *testing.T, expected string, actual string) {
-	if runtime.GOOS == "windows" {
-		expected = expected + "\r\n"
-	} else {
-		expected = expected + "\n"
-	}
-
-	assert.Equal(t, expected, actual)
-}

multiplexed_writer.go

@@ -0,0 +1,29 @@
+package cmd
+
+import (
+	"fmt"
+	"io"
+)
+
+// NewMultiplexedWriter returns a new multiplexer
+func NewMultiplexedWriter(outputs ...io.Writer) *MultiplexedWriter {
+	return &MultiplexedWriter{Outputs: outputs}
+}
+
+// MultiplexedWriter writes to multiple writers at once
+type MultiplexedWriter struct {
+	Outputs []io.Writer
+}
+
+// Write writes the given bytes. If one write fails it returns the error
+// and bytes of the failed write operation
+func (w MultiplexedWriter) Write(p []byte) (n int, err error) {
+	for _, o := range w.Outputs {
+		n, err = o.Write(p)
+		if err != nil {
+			return 0, fmt.Errorf("Error in writer: %s", err.Error())
+		}
+	}
+
+	return n, nil
+}

multiplexed_writer_test.go

@@ -0,0 +1,57 @@
+package cmd
+
+import (
+	"bytes"
+	"fmt"
+	"github.com/stretchr/testify/assert"
+	"os"
+	"testing"
+)
+
+func TestMultiplexedWriter(t *testing.T) {
+	writer01 := bytes.Buffer{}
+	writer02 := bytes.Buffer{}
+	// Test another io.Writer interface type
+	r, w, _ := os.Pipe()
+
+	writer := NewMultiplexedWriter(&writer01, &writer02, w)
+	n, err := writer.Write([]byte(`test`))
+
+	assert.Nil(t, err)
+	assert.Equal(t, 4, n)
+	assert.Equal(t, "test", writer01.String())
+	assert.Equal(t, "test", writer02.String())
+
+	data := make([]byte, 4)
+	n, err = r.Read(data)
+	assert.Nil(t, err)
+	assert.Equal(t, 4, n)
+	assert.Equal(t, "test", string(data))
+}
+
+func TestMultiplexedWriter_SingleWirter(t *testing.T) {
+	writer01 := bytes.Buffer{}
+
+	writer := NewMultiplexedWriter(&writer01)
+
+	n, _ := writer.Write([]byte(`another`))
+
+	assert.Equal(t, 7, n)
+	assert.Equal(t, "another", writer01.String())
+}
+
+func TestMultiplexedWriter_Fail(t *testing.T) {
+	writer := NewMultiplexedWriter(InvalidWriter{})
+
+	n, err := writer.Write([]byte(`another`))
+
+	assert.Equal(t, 0, n)
+	assert.Equal(t, "Error in writer: failed", err.Error())
+}
+
+type InvalidWriter struct {
+}
+
+func (w InvalidWriter) Write(p []byte) (n int, err error) {
+	return 0, fmt.Errorf("failed")
+}

testing_utils.go

@@ -0,0 +1,55 @@
+package cmd
+
+import (
+	"bytes"
+	"github.com/stretchr/testify/assert"
+	"io"
+	"log"
+	"os"
+	"runtime"
+	"sync"
+	"testing"
+)
+
+// CaptureStandardOutput allows to capture the output which will be written
+// to os.Stdout and os.Stderr.
+// It returns the captured output and the return value of the called function
+func CaptureStandardOutput(f func() interface{}) (string, interface{}) {
+	reader, writer, err := os.Pipe()
+	if err != nil {
+		panic(err)
+	}
+	stdout := os.Stdout
+	stderr := os.Stderr
+	defer func() {
+		os.Stdout = stdout
+		os.Stderr = stderr
+		log.SetOutput(os.Stderr)
+	}()
+	os.Stdout = writer
+	os.Stderr = writer
+	log.SetOutput(writer)
+	out := make(chan string)
+	wg := new(sync.WaitGroup)
+	wg.Add(1)
+	go func() {
+		var buf bytes.Buffer
+		wg.Done()
+		io.Copy(&buf, reader)
+		out <- buf.String()
+	}()
+	wg.Wait()
+	result := f()
+	writer.Close()
+	return <-out, result
+}
+
+func assertEqualWithLineBreak(t *testing.T, expected string, actual string) {
+	if runtime.GOOS == "windows" {
+		expected = expected + "\r\n"
+	} else {
+		expected = expected + "\n"
+	}
+
+	assert.Equal(t, expected, actual)
+}