diff --git a/shell.go b/shell.go index 56ed071..11d7a32 100644 --- a/shell.go +++ b/shell.go @@ -58,11 +58,12 @@ 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() } @@ -70,7 +71,7 @@ func (c *Command) ProcFn() func(...interface{}) *Process { 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 } diff --git a/shell_test.go b/shell_test.go index a1606b8..2f77213 100644 --- a/shell_test.go +++ b/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)) + } +}