Enable per-process Tee for access to proc Stdout and Stderr streams

shell.go

@@ -58,19 +58,20 @@ type Command struct {
 	args []string
 	in   *Command
 	wd   string
+	tee  io.Writer
 }
 
 func (c *Command) ProcFn() func(...interface{}) *Process {
 	return func(args ...interface{}) *Process {
-		cmd := &Command{c.args, c.in, c.wd}
+		cmd := &Command{c.args, c.in, c.wd, c.tee}
 		cmd.addArgs(args...)
 		return cmd.Run()
 	}
 }
 
 func (c *Command) OutputFn() func(...interface{}) (string, error) {
 	return func(args ...interface{}) (out string, err error) {
-		cmd := &Command{c.args, c.in, c.wd}
+		cmd := &Command{c.args, c.in, c.wd, c.tee}
 		cmd.addArgs(args...)
 		defer func() {
 			if p, ok := recover().(*Process); p != nil {
@@ -88,7 +89,7 @@ func (c *Command) OutputFn() func(...interface{}) (string, error) {
 
 func (c *Command) ErrFn() func(...interface{}) error {
 	return func(args ...interface{}) (err error) {
-		cmd := &Command{c.args, c.in, c.wd}
+		cmd := &Command{c.args, c.in, c.wd, c.tee}
 		cmd.addArgs(args...)
 		defer func() {
 			if p, ok := recover().(*Process); p != nil {
@@ -133,6 +134,11 @@ func (c *Command) addArgs(args ...interface{}) {
 	c.args = append(c.args, strArgs...)
 }
 
+func (c *Command) Tee(t io.Writer) *Command {
+	c.tee = t
+	return c
+}
+
 func (c *Command) shellCmd(quote bool) string {
 	if !quote {
 		return strings.Join(c.args, " ")
@@ -169,16 +175,19 @@ func (c *Command) execute(cmd *exec.Cmd, call func() error) *Process {
 		assert(err)
 		p.Stdin = stdin
 	}
-	var stdout bytes.Buffer
 	if Tee != nil {
-		cmd.Stdout = io.MultiWriter(&stdout, Tee)
+		c.tee = Tee
+	}
+	var stdout bytes.Buffer
+	if c.tee != nil {
+		cmd.Stdout = io.MultiWriter(&stdout, c.tee)
 	} else {
 		cmd.Stdout = &stdout
 	}
 	p.Stdout = &stdout
 	var stderr bytes.Buffer
-	if Tee != nil {
-		cmd.Stderr = io.MultiWriter(&stderr, Tee)
+	if c.tee != nil {
+		cmd.Stderr = io.MultiWriter(&stderr, c.tee)
 	} else {
 		cmd.Stderr = &stderr
 	}

shell_test.go

@@ -3,6 +3,8 @@ package shell
 import (
 	"bytes"
 	"fmt"
+	"io"
+	"io/ioutil"
 	"os"
 	"strings"
 	"testing"
@@ -182,3 +184,28 @@ func TestSetWorkDir(t *testing.T) {
 		t.Errorf("expected touched file to be present in correct working dir but was not")
 	}
 }
+
+func TestCmdTee(t *testing.T) {
+	pr, pw := io.Pipe()
+	go func() {
+		defer pw.Close()
+		p := Cmd("echo", "test").Tee(pw).Run()
+		if p.ExitStatus != 0 {
+			t.Fatal(p.Error())
+		}
+
+		if p.String() != "test" {
+			t.Errorf("expected String() output to be (test), but was (%s)", string(p.String()))
+		}
+	}()
+
+	out, err := ioutil.ReadAll(pr)
+	if err != nil {
+		t.Fatal(err)
+	}
+	pr.Close()
+
+	if string(out) != "test\n" {
+		t.Errorf("expected Tee output to be (test\\n), but was (%s)", string(out))
+	}
+}