testscript: add waitmatch command

testscript/cmd.go

@@ -9,7 +9,6 @@ import (
 	"bytes"
 	"fmt"
 	"os"
-	"os/exec"
 	"path/filepath"
 	"regexp"
 	"runtime"
@@ -27,30 +26,31 @@ import (
 //
 // NOTE: If you make changes here, update doc.go.
 var scriptCmds = map[string]func(*TestScript, bool, []string){
-	"cd":       (*TestScript).cmdCd,
-	"chmod":    (*TestScript).cmdChmod,
-	"cmp":      (*TestScript).cmdCmp,
-	"cmpenv":   (*TestScript).cmdCmpenv,
-	"cp":       (*TestScript).cmdCp,
-	"env":      (*TestScript).cmdEnv,
-	"exec":     (*TestScript).cmdExec,
-	"exists":   (*TestScript).cmdExists,
-	"grep":     (*TestScript).cmdGrep,
-	"kill":     (*TestScript).cmdKill,
-	"mkdir":    (*TestScript).cmdMkdir,
-	"mv":       (*TestScript).cmdMv,
-	"rm":       (*TestScript).cmdRm,
-	"skip":     (*TestScript).cmdSkip,
-	"stderr":   (*TestScript).cmdStderr,
-	"stdin":    (*TestScript).cmdStdin,
-	"stdout":   (*TestScript).cmdStdout,
-	"ttyin":    (*TestScript).cmdTtyin,
-	"ttyout":   (*TestScript).cmdTtyout,
-	"stop":     (*TestScript).cmdStop,
-	"symlink":  (*TestScript).cmdSymlink,
-	"unix2dos": (*TestScript).cmdUNIX2DOS,
-	"unquote":  (*TestScript).cmdUnquote,
-	"wait":     (*TestScript).cmdWait,
+	"cd":        (*TestScript).cmdCd,
+	"chmod":     (*TestScript).cmdChmod,
+	"cmp":       (*TestScript).cmdCmp,
+	"cmpenv":    (*TestScript).cmdCmpenv,
+	"cp":        (*TestScript).cmdCp,
+	"env":       (*TestScript).cmdEnv,
+	"exec":      (*TestScript).cmdExec,
+	"exists":    (*TestScript).cmdExists,
+	"grep":      (*TestScript).cmdGrep,
+	"kill":      (*TestScript).cmdKill,
+	"mkdir":     (*TestScript).cmdMkdir,
+	"mv":        (*TestScript).cmdMv,
+	"rm":        (*TestScript).cmdRm,
+	"skip":      (*TestScript).cmdSkip,
+	"stderr":    (*TestScript).cmdStderr,
+	"stdin":     (*TestScript).cmdStdin,
+	"stdout":    (*TestScript).cmdStdout,
+	"ttyin":     (*TestScript).cmdTtyin,
+	"ttyout":    (*TestScript).cmdTtyout,
+	"stop":      (*TestScript).cmdStop,
+	"symlink":   (*TestScript).cmdSymlink,
+	"unix2dos":  (*TestScript).cmdUNIX2DOS,
+	"unquote":   (*TestScript).cmdUnquote,
+	"wait":      (*TestScript).cmdWait,
+	"waitmatch": (*TestScript).cmdWaitMatch,
 }
 
 // cd changes to a different directory.
@@ -242,15 +242,10 @@ func (ts *TestScript) cmdExec(neg bool, args []string) {
 		if ts.findBackground(bgName) != nil {
 			ts.Fatalf("duplicate background process name %q", bgName)
 		}
-		var cmd *exec.Cmd
-		cmd, err = ts.execBackground(args[0], args[1:len(args)-1]...)
+		var bg backgroundCmd
+		bg, err = ts.execBackground(bgName, neg, args[0], args[1:len(args)-1]...)
 		if err == nil {
-			wait := make(chan struct{})
-			go func() {
-				waitOrStop(ts.ctxt, cmd, -1)
-				close(wait)
-			}()
-			ts.background = append(ts.background, backgroundCmd{bgName, cmd, wait, neg})
+			ts.background = append(ts.background, bg)
 		}
 		ts.stdout, ts.stderr = "", ""
 	} else {
@@ -568,14 +563,57 @@ func (ts *TestScript) cmdWait(neg bool, args []string) {
 	}
 }
 
+// waitmatch waits for a background command to print a line matching a pattern.
+func (ts *TestScript) cmdWaitMatch(neg bool, args []string) {
+	if neg {
+		ts.Fatalf("unsupported: ! waitmatch")
+	}
+	useStderr := len(args) > 0 && args[0] == "-stderr"
+	if useStderr {
+		args = args[1:]
+	}
+	if len(args) < 2 {
+		ts.Fatalf("usage: waitmatch [-stderr] command pattern [var...]")
+	}
+	bgName, pattern, vars := args[0], args[1], args[2:]
+	bg := ts.findBackground(bgName)
+	if bg == nil {
+		ts.Fatalf("unknown background process %q", bgName)
+	}
+	re, err := regexp.Compile(pattern)
+	ts.Check(err)
+	if n := re.NumSubexp(); len(vars) > n {
+		ts.Fatalf("cannot set %d variables from %d subexpressions", len(vars), n)
+	}
+	out, stream := bg.stdout, "stdout"
+	if useStderr {
+		out, stream = bg.stderr, "stderr"
+	}
+	want := fmt.Sprintf("a line matching %#q to %s", pattern, stream)
+	scanner := bufio.NewScanner(out.reader(ts.ctxt))
+	for scanner.Scan() {
+		if m := re.FindStringSubmatch(scanner.Text()); m != nil {
+			for i, name := range vars {
+				ts.Setenv(name, m[1+i])
+			}
+			return
+		}
+	}
+	if ts.ctxt.Err() != nil {
+		ts.Fatalf("test timed out while waiting for %q to print %s", bgName, want)
+	}
+	ts.Check(scanner.Err())
+	ts.Fatalf("%q exited without printing %s", bgName, want)
+}
+
 func (ts *TestScript) waitBackgroundOne(bgName string) {
 	bg := ts.findBackground(bgName)
 	if bg == nil {
 		ts.Fatalf("unknown background process %q", bgName)
 	}
 	<-bg.wait
-	ts.stdout = bg.cmd.Stdout.(*strings.Builder).String()
-	ts.stderr = bg.cmd.Stderr.(*strings.Builder).String()
+	ts.stdout = bg.stdout.String()
+	ts.stderr = bg.stderr.String()
 	if ts.stdout != "" {
 		fmt.Fprintf(&ts.log, "[stdout]\n%s", ts.stdout)
 	}
@@ -625,13 +663,13 @@ func (ts *TestScript) waitBackground(checkStatus bool) {
 		args := append([]string{filepath.Base(bg.cmd.Args[0])}, bg.cmd.Args[1:]...)
 		fmt.Fprintf(&ts.log, "[background] %s: %v\n", strings.Join(args, " "), bg.cmd.ProcessState)
 
-		cmdStdout := bg.cmd.Stdout.(*strings.Builder).String()
+		cmdStdout := bg.stdout.String()
 		if cmdStdout != "" {
 			fmt.Fprintf(&ts.log, "[stdout]\n%s", cmdStdout)
 			stdouts = append(stdouts, cmdStdout)
 		}
 
-		cmdStderr := bg.cmd.Stderr.(*strings.Builder).String()
+		cmdStderr := bg.stderr.String()
 		if cmdStderr != "" {
 			fmt.Fprintf(&ts.log, "[stderr]\n%s", cmdStderr)
 			stderrs = append(stderrs, cmdStderr)

testscript/doc.go

@@ -248,6 +248,16 @@ The predefined commands are:
     must have been started with the final token '&command&` as described for the
     exec command.
 
+  - waitmatch [-stderr] command pattern [var...]
+    Wait for a command started in the background with the final token
+    '&command&' to print a line matching the regular expression pattern.
+    It matches against the command's standard output unless -stderr is
+    specified. This is useful to wait for a background process to be ready,
+    such as a server printing the address it listens on. Any variables given
+    are set to the subexpressions of the match, in order. It fails if the
+    command exits without printing a matching line, or if the test times out
+    first.
+
 When TestScript runs a script and the script fails, by default TestScript shows
 the execution of the most recent phase of the script (since the last # comment)
 and only shows the # comments for earlier phases. For example, here is a

testscript/testdata/interrupt.txt

@@ -1,7 +1,7 @@
 [windows] skip
 
-signalcatcher &
-waitfile catchsignal
+signalcatcher &sc&
+waitmatch sc ^ready$
 interrupt
 wait
 stdout 'caught interrupt'

testscript/testdata/interrupt_implicit.txt

@@ -1,4 +1,4 @@
 # Let testscript stop signalcatcher at the end of the testscript.
 
-signalcatcher &
-waitfile catchsignal
+signalcatcher &sc&
+waitmatch sc ^ready$

testscript/testdata/waitmatch.txt

@@ -0,0 +1,53 @@
+# waitmatch waits for a background command to print a matching line,
+# including one printed before waitmatch starts waiting,
+# and sets variables to the subexpressions of the match.
+! exec printandwait stdout 'starting up' 'listening on 127.0.0.1:1234' 'more output' &srv&
+waitmatch srv '^listening on (.+):(\d+)$' HOST PORT
+exec printargs $HOST $PORT
+stdout '"127.0.0.1" "1234"'
+waitmatch srv '^starting up$'
+waitmatch srv '^more output$'
+
+# Waiting for the command still gives all of its output.
+kill srv
+wait srv
+stdout '^starting up\nlistening on 127.0.0.1:1234\nmore output\n$'
+
+# -stderr waits for a line on standard error instead.
+! exec printandwait stderr 'ready' &errsrv&
+waitmatch -stderr errsrv '^ready$'
+
+# A command may also exit before or while waitmatch waits.
+exec fprintargs stdout done &short&
+waitmatch short '^done$'
+
+! testscript -v exited
+stdout 'FAIL: .*: "bg" exited without printing a line matching `never` to stdout'
+
+! testscript -timeout=500ms timeout
+stdout 'FAIL: .*: test timed out while waiting for "bg" to print a line matching `never` to stdout'
+
+! testscript usage
+stdout 'unknown background process "nosuch"'
+stdout 'cannot set 2 variables from 1 subexpressions'
+
+# Lines may also end with CRLF.
+[!exec:cat] stop
+unix2dos crlf.txt
+exec cat crlf.txt &crlf&
+waitmatch crlf '^ready$'
+
+-- exited/script.txt --
+exec fprintargs stdout foo &bg&
+waitmatch bg never
+-- timeout/script.txt --
+exec printandwait stdout foo &bg&
+waitmatch bg never
+-- usage/script.txt --
+exec printandwait stdout foo &bg&
+waitmatch nosuch foo
+-- usage/twovars.txt --
+exec printandwait stdout foo &bg&
+waitmatch bg '(foo)' A B
+-- crlf.txt --
+ready

testscript/testscript.go

@@ -24,6 +24,7 @@ import (
 	"slices"
 	"strconv"
 	"strings"
+	"sync"
 	"sync/atomic"
 	"syscall"
 	"testing"
@@ -424,10 +425,12 @@ type TestScript struct {
 }
 
 type backgroundCmd struct {
-	name string
-	cmd  *exec.Cmd
-	wait <-chan struct{}
-	neg  bool // if true, cmd should fail
+	name   string
+	cmd    *exec.Cmd
+	stdout *backgroundOutput
+	stderr *backgroundOutput
+	wait   <-chan struct{}
+	neg    bool // if true, cmd should fail
 }
 
 func writeFile(name string, data []byte, perm fs.FileMode, excl bool) error {
@@ -1084,23 +1087,112 @@ func (ts *TestScript) exec(command string, args ...string) (stdout, stderr strin
 }
 
 // execBackground starts the given command line (an actual subprocess, not simulated)
-// in ts.cd with environment ts.env.
-func (ts *TestScript) execBackground(command string, args ...string) (*exec.Cmd, error) {
+// in ts.cd with environment ts.env, as a background command with the given name.
+func (ts *TestScript) execBackground(name string, neg bool, command string, args ...string) (backgroundCmd, error) {
 	if ts.ttyin != "" {
-		return nil, errors.New("ttyin is not supported by background commands")
+		return backgroundCmd{}, errors.New("ttyin is not supported by background commands")
 	}
 	cmd, err := ts.buildExecCmd(command, args...)
 	if err != nil {
-		return nil, err
+		return backgroundCmd{}, err
+	}
+	wait := make(chan struct{})
+	bg := backgroundCmd{
+		name:   name,
+		cmd:    cmd,
+		stdout: newBackgroundOutput(),
+		stderr: newBackgroundOutput(),
+		wait:   wait,
+		neg:    neg,
 	}
 	cmd.Dir = ts.cd
 	cmd.Env = append(ts.env, "PWD="+ts.cd)
-	var stdoutBuf, stderrBuf strings.Builder
 	cmd.Stdin = strings.NewReader(ts.stdin)
-	cmd.Stdout = &stdoutBuf
-	cmd.Stderr = &stderrBuf
+	cmd.Stdout = bg.stdout
+	cmd.Stderr = bg.stderr
 	ts.stdin = ""
-	return cmd, cmd.Start()
+	if err := cmd.Start(); err != nil {
+		return backgroundCmd{}, err
+	}
+	go func() {
+		waitOrStop(ts.ctxt, cmd, -1)
+		bg.stdout.close()
+		bg.stderr.close()
+		close(wait)
+	}()
+	return bg, nil
+}
+
+// backgroundOutput buffers an output stream of a background command.
+// It is safe for concurrent use, so that waitmatch can read the output
+// while the command is still running.
+type backgroundOutput struct {
+	mu      sync.Mutex
+	buf     strings.Builder
+	written chan struct{} // closed by the next write or by close
+	closed  bool
+}
+
+func newBackgroundOutput() *backgroundOutput {
+	return &backgroundOutput{written: make(chan struct{})}
+}
+
+func (o *backgroundOutput) Write(p []byte) (int, error) {
+	o.mu.Lock()
+	defer o.mu.Unlock()
+	o.buf.Write(p)
+	close(o.written)
+	o.written = make(chan struct{})
+	return len(p), nil
+}
+
+// close marks the end of the output, once the command has exited.
+func (o *backgroundOutput) close() {
+	o.mu.Lock()
+	defer o.mu.Unlock()
+	o.closed = true
+	close(o.written)
+}
+
+// String returns the output written so far.
+func (o *backgroundOutput) String() string {
+	o.mu.Lock()
+	defer o.mu.Unlock()
+	return o.buf.String()
+}
+
+// reader returns a reader which produces the output from the start,
+// blocking until more is written. It returns [io.EOF] once the output is closed,
+// or ctx.Err() if the context is cancelled first.
+func (o *backgroundOutput) reader(ctx context.Context) io.Reader {
+	return &backgroundOutputReader{o: o, ctx: ctx}
+}
+
+type backgroundOutputReader struct {
+	o   *backgroundOutput
+	ctx context.Context
+	off int
+}
+
+func (r *backgroundOutputReader) Read(p []byte) (int, error) {
+	for {
+		r.o.mu.Lock()
+		s, written, closed := r.o.buf.String(), r.o.written, r.o.closed
+		r.o.mu.Unlock()
+		if r.off < len(s) {
+			n := copy(p, s[r.off:])
+			r.off += n
+			return n, nil
+		}
+		if closed {
+			return 0, io.EOF
+		}
+		select {
+		case <-written:
+		case <-r.ctx.Done():
+			return 0, r.ctx.Err()
+		}
+	}
 }
 
 func (ts *TestScript) buildExecCmd(command string, args ...string) (*exec.Cmd, error) {