testscript: expose (*TestScript).stdout via Stdout()

testscript/testdata/cmd_stdout_stderr.txt

@@ -0,0 +1,39 @@
+# Verify that when we don't update stdout when we don't attempt to write via Stdout()
+fprintargs stdout hello stdout from fprintargs
+stdout 'hello stdout from fprintargs'
+echoandexit 0
+stdout 'hello stdout from fprintargs'
+
+# Verify that when we don't update stderr when we don't attempt to write via Stderr()
+fprintargs stderr hello stderr from fprintargs
+stderr 'hello stderr from fprintargs'
+echoandexit 0
+stderr 'hello stderr from fprintargs'
+
+# Verify that we do update stdout when we attempt to write via Stdout() or Stderr()
+fprintargs stdout hello stdout from fprintargs
+stdout 'hello stdout from fprintargs'
+! stderr .+
+echoandexit 0 'hello stdout from echoandexit'
+stdout 'hello stdout from echoandexit'
+! stderr .+
+fprintargs stdout hello stdout from fprintargs
+stdout 'hello stdout from fprintargs'
+! stderr .+
+echoandexit 0 '' 'hello stderr from echoandexit'
+! stdout .+
+stderr 'hello stderr from echoandexit'
+
+# Verify that we do update stderr when we attempt to write via Stdout() or Stderr()
+fprintargs stderr hello stderr from fprintargs
+! stdout .+
+stderr 'hello stderr from fprintargs'
+echoandexit 0 'hello stdout from echoandexit'
+stdout 'hello stdout from echoandexit'
+! stderr .+
+fprintargs stdout hello stdout from fprintargs
+stdout 'hello stdout from fprintargs'
+! stderr .+
+echoandexit 0 '' 'hello stderr from echoandexit'
+! stdout .+
+stderr 'hello stderr from echoandexit'

testscript/testdata/testscript_stdout_stderr_error.txt

@@ -0,0 +1,20 @@
+# Verify that stdout and stderr get set event when a user-builtin
+# command aborts. Note that we need to assert against stdout
+# because our meta testscript command sees only a single log.
+unquote scripts/testscript.txt
+! testscript -v scripts
+cmpenv stdout stdout.golden
+
+-- scripts/testscript.txt --
+> printargs hello world
+> echoandexit 1 'this is stdout' 'this is stderr'
+-- stdout.golden --
+>  printargs hello world
+[stdout]
+["printargs" "hello" "world"]
+>  echoandexit 1 'this is stdout' 'this is stderr'
+[stdout]
+this is stdout
+[stderr]
+this is stderr
+FAIL: ${$}WORK${/}scripts${/}testscript.txt:2: told to exit with code 1

testscript/testscript.go

@@ -14,6 +14,7 @@ import (
 	"flag"
 	"fmt"
 	"go/build"
+	"io"
 	"io/fs"
 	"io/ioutil"
 	"os"
@@ -362,6 +363,17 @@ type TestScript struct {
 	scriptFiles   map[string]string           // files stored in the txtar archive (absolute paths -> path in script)
 	scriptUpdates map[string]string           // updates to testscript files via UpdateScripts.
 
+	// runningBuiltin indicates if we are running a user-supplied builtin
+	// command. These commands are specified via Params.Cmds.
+	runningBuiltin bool
+
+	// builtinStd(out|err) are established if a user-supplied builtin command
+	// requests Stdout() or Stderr(). Either both are non-nil, or both are nil.
+	// This invariant is maintained by both setBuiltinStd() and
+	// clearBuiltinStd().
+	builtinStdout *strings.Builder
+	builtinStderr *strings.Builder
+
 	ctxt        context.Context // per TestScript context
 	gracePeriod time.Duration   // time between SIGQUIT and SIGKILL
 }
@@ -659,10 +671,29 @@ func (ts *TestScript) runLine(line string) (runOK bool) {
 	if cmd == nil {
 		ts.Fatalf("unknown command %q", args[0])
 	}
-	cmd(ts, neg, args[1:])
+	ts.callBuiltinCmd(args[0], func() {
+		cmd(ts, neg, args[1:])
+	})
 	return true
 }
 
+func (ts *TestScript) callBuiltinCmd(cmd string, runCmd func()) {
+	ts.runningBuiltin = true
+	defer func() {
+		r := recover()
+		ts.runningBuiltin = false
+		ts.clearBuiltinStd()
+		switch r {
+		case nil:
+			// we did not panic
+		default:
+			// re-"throw" the panic
+			panic(r)
+		}
+	}()
+	runCmd()
+}
+
 func (ts *TestScript) applyScriptUpdates() {
 	if len(ts.scriptUpdates) == 0 {
 		return
@@ -788,6 +819,60 @@ func (ts *TestScript) Check(err error) {
 	}
 }
 
+// Stdout returns an io.Writer that can be used by a user-supplied builtin
+// command (delcared via Params.Cmds) to write to stdout. If this method is
+// called outside of the execution of a user-supplied builtin command, the
+// call panics.
+func (ts *TestScript) Stdout() io.Writer {
+	if !ts.runningBuiltin {
+		panic("can only call TestScript.Stdout when running a builtin command")
+	}
+	ts.setBuiltinStd()
+	return ts.builtinStdout
+}
+
+// Stderr returns an io.Writer that can be used by a user-supplied builtin
+// command (delcared via Params.Cmds) to write to stderr. If this method is
+// called outside of the execution of a user-supplied builtin command, the
+// call panics.
+func (ts *TestScript) Stderr() io.Writer {
+	if !ts.runningBuiltin {
+		panic("can only call TestScript.Stderr when running a builtin command")
+	}
+	ts.setBuiltinStd()
+	return ts.builtinStderr
+}
+
+// setBuiltinStd ensures that builtinStdout and builtinStderr are non nil.
+func (ts *TestScript) setBuiltinStd() {
+	// This method must maintain the invariant that both builtinStdout and
+	// builtinStderr are set or neither are set
+
+	// If both are set, nothing to do
+	if ts.builtinStdout != nil && ts.builtinStderr != nil {
+		return
+	}
+	ts.builtinStdout = new(strings.Builder)
+	ts.builtinStderr = new(strings.Builder)
+}
+
+// clearBuiltinStd sets ts.stdout and ts.stderr from the builtin command
+// buffers, logs both, and resets both builtinStdout and builtinStderr to nil.
+func (ts *TestScript) clearBuiltinStd() {
+	// This method must maintain the invariant that both builtinStdout and
+	// builtinStderr are set or neither are set
+
+	// If neither set, nothing to do
+	if ts.builtinStdout == nil && ts.builtinStderr == nil {
+		return
+	}
+	ts.stdout = ts.builtinStdout.String()
+	ts.builtinStdout = nil
+	ts.stderr = ts.builtinStderr.String()
+	ts.builtinStderr = nil
+	ts.logStd()
+}
+
 // Logf appends the given formatted message to the test log transcript.
 func (ts *TestScript) Logf(format string, args ...interface{}) {
 	format = strings.TrimSuffix(format, "\n")
@@ -933,13 +1018,18 @@ func interruptProcess(p *os.Process) {
 func (ts *TestScript) Exec(command string, args ...string) error {
 	var err error
 	ts.stdout, ts.stderr, err = ts.exec(command, args...)
+	ts.logStd()
+	return err
+}
+
+// logStd logs the current non-empty values of stdout and stderr.
+func (ts *TestScript) logStd() {
 	if ts.stdout != "" {
 		ts.Logf("[stdout]\n%s", ts.stdout)
 	}
 	if ts.stderr != "" {
 		ts.Logf("[stderr]\n%s", ts.stderr)
 	}
-	return err
 }
 
 // expand applies environment variable expansion to the string s.
@@ -954,6 +1044,12 @@ func (ts *TestScript) expand(s string) string {
 
 // fatalf aborts the test with the given failure message.
 func (ts *TestScript) Fatalf(format string, args ...interface{}) {
+	// In user-supplied builtins, the only way we have of aborting
+	// is via Fatalf. Hence if we are aborting from a user-supplied
+	// builtin, it's important we first log stdout and stderr. If
+	// we are not, the following call is a no-op.
+	ts.clearBuiltinStd()
+
 	fmt.Fprintf(&ts.log, "FAIL: %s:%d: %s\n", ts.file, ts.lineno, fmt.Sprintf(format, args...))
 	// This should be caught by the defer inside the TestScript.runLine method.
 	// We do this rather than calling ts.t.FailNow directly because we want to

testscript/testscript_test.go

@@ -234,11 +234,14 @@ func TestScripts(t *testing.T) {
 						Cmds: map[string]func(ts *TestScript, neg bool, args []string){
 							"some-param-cmd": func(ts *TestScript, neg bool, args []string) {
 							},
+							"echoandexit": echoandexit,
 						},
 						ContinueOnError: *fContinue,
 					})
 				}()
-				ts.stdout = strings.Replace(t.log.String(), ts.workdir, "$WORK", -1)
+				stdout := t.log.String()
+				stdout = strings.ReplaceAll(stdout, ts.workdir, "$WORK")
+				fmt.Fprint(ts.Stdout(), stdout)
 				if neg {
 					if !t.failed {
 						ts.Fatalf("testscript unexpectedly succeeded")
@@ -249,6 +252,7 @@ func TestScripts(t *testing.T) {
 					ts.Fatalf("testscript unexpectedly failed with errors: %q", &t.log)
 				}
 			},
+			"echoandexit": echoandexit,
 		},
 		Setup: func(env *Env) error {
 			infos, err := ioutil.ReadDir(env.WorkDir)
@@ -274,6 +278,33 @@ func TestScripts(t *testing.T) {
 	// TODO check that the temp directory has been removed.
 }
 
+func echoandexit(ts *TestScript, neg bool, args []string) {
+	// Takes at least one argument
+	//
+	// args[0] - int that indicates the exit code of the command
+	// args[1] - the string to echo to stdout if non-empty
+	// args[2] - the string to echo to stderr if non-empty
+	if len(args) == 0 || len(args) > 3 {
+		ts.Fatalf("echoandexit takes at least one and at most three arguments")
+	}
+	if neg {
+		ts.Fatalf("neg means nothing for echoandexit")
+	}
+	exitCode, err := strconv.ParseInt(args[0], 10, 64)
+	if err != nil {
+		ts.Fatalf("failed to parse exit code from %q: %v", args[0], err)
+	}
+	if len(args) > 1 && args[1] != "" {
+		fmt.Fprint(ts.Stdout(), args[1])
+	}
+	if len(args) > 2 && args[2] != "" {
+		fmt.Fprint(ts.Stderr(), args[2])
+	}
+	if exitCode != 0 {
+		ts.Fatalf("told to exit with code %d", exitCode)
+	}
+}
+
 // TestTestwork tests that using the flag -testwork will make sure the work dir isn't removed
 // after the test is done. It uses an empty testscript file that doesn't do anything.
 func TestTestwork(t *testing.T) {