testscript: add waitmatch command

It blocks until a background process prints a matching line to stdout.
Any given env var names are set to the subexpression strings as well.

This is useful to wait for a background process to be ready before
running more commands to interact with it, such as waiting for an HTTP
server to listen on a TCP port before sending it the first request.

In fact, our interrupt tests already had this problem; we only worked
around it by having the process create a file when it's ready,
and adding a custom command which would wait for the file by polling.
It's much nicer and faster to use waitmatch instead,
and as a bonus, we can refactor existing tests rather than adding more.

Author: mvdan

testscript/cmd.go

@@ -8,6 +8,7 @@ import (
 	"bufio"
 	"bytes"
 	"fmt"
+	"io"
 	"os"
 	"os/exec"
 	"path/filepath"
@@ -26,30 +27,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.
@@ -237,14 +239,26 @@ func (ts *TestScript) cmdExec(neg bool, args []string) {
 			ts.Fatalf("duplicate background process name %q", bgName)
 		}
 		var cmd *exec.Cmd
-		cmd, err = ts.execBackground(args[0], args[1:len(args)-1]...)
+		var outreader io.ReadCloser
+		cmd, outreader, err = ts.execBackground(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})
+			outbuf := new(strings.Builder)
+			ts.background = append(ts.background, backgroundCmd{
+				name:         bgName,
+				cmd:          cmd,
+				stdoutBuffer: outbuf,
+				stdoutReader: struct {
+					io.Reader
+					io.Closer
+				}{io.TeeReader(outreader, outbuf), outreader},
+				waitc: wait,
+				neg:   neg,
+			})
 		}
 		ts.stdout, ts.stderr = "", ""
 	} else {
@@ -567,9 +581,7 @@ func (ts *TestScript) waitBackgroundOne(bgName string) {
 	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, ts.stderr = bg.wait()
 	if ts.stdout != "" {
 		fmt.Fprintf(&ts.log, "[stdout]\n%s", ts.stdout)
 	}
@@ -614,18 +626,15 @@ func (ts *TestScript) findBackground(bgName string) *backgroundCmd {
 func (ts *TestScript) waitBackground(checkStatus bool) {
 	var stdouts, stderrs []string
 	for _, bg := range ts.background {
-		<-bg.wait
+		cmdStdout, cmdStderr := bg.wait()
 
 		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()
 		if cmdStdout != "" {
 			fmt.Fprintf(&ts.log, "[stdout]\n%s", cmdStdout)
 			stdouts = append(stdouts, cmdStdout)
 		}
-
-		cmdStderr := bg.cmd.Stderr.(*strings.Builder).String()
 		if cmdStderr != "" {
 			fmt.Fprintf(&ts.log, "[stderr]\n%s", cmdStderr)
 			stderrs = append(stderrs, cmdStderr)
@@ -652,6 +661,69 @@ func (ts *TestScript) waitBackground(checkStatus bool) {
 	ts.background = nil
 }
 
+// cmdWaitMatch waits until a background command prints a line to standard output
+// which matches the given regular expression. Once a match is found, the given
+// environment variable names are set to the subexpressions of the match.
+func (ts *TestScript) cmdWaitMatch(neg bool, args []string) {
+	if len(args) < 1 {
+		ts.Fatalf("usage: waitmatch name regexp [env-var...]")
+	}
+	if neg {
+		ts.Fatalf("unsupported: ! waitmatch")
+	}
+	bg := ts.findBackground(args[0])
+	if bg == nil {
+		ts.Fatalf("unknown background process %q", args[0])
+	}
+	rx, err := regexp.Compile(args[1])
+	ts.Check(err)
+	envs := args[2:]
+	if n := rx.NumSubexp(); n < len(envs) {
+		ts.Fatalf("cannot extract %d subexpressions into %d env vars", n, len(envs))
+	}
+	for {
+		line, err := readLine(bg.stdoutReader)
+		ts.Check(err)
+		m := rx.FindSubmatch(line)
+		if m != nil {
+			subm := m[1:]
+			for i, env := range envs {
+				ts.Setenv(env, string(subm[i]))
+			}
+		}
+		if err == io.EOF {
+			if m == nil {
+				ts.Fatalf("reached EOF without matching any line")
+			}
+			return
+		} else {
+			ts.Check(err)
+		}
+		if m != nil {
+			break
+		}
+	}
+}
+
+// readLine consumes enough bytes to read a line.
+func readLine(r io.Reader) ([]byte, error) {
+	var line []byte
+	for {
+		var buf [1]byte
+		n, err := r.Read(buf[:])
+		if n > 0 {
+			b := buf[0]
+			if b == '\n' {
+				return line, nil
+			}
+			line = append(line, b)
+		}
+		if err != nil {
+			return line, err
+		}
+	}
+}
+
 // scriptMatch implements both stdout and stderr.
 func scriptMatch(ts *TestScript, neg bool, args []string, text, name string) {
 	n := 0

testscript/testdata/interrupt.txt

@@ -1,7 +1,20 @@
 [windows] skip
 
-signalcatcher &
-waitfile catchsignal
+# Start a background process, wait for it to be ready by printing a line,
+# send it an interrupt to stop, and wait for it to stop.
+
+signalcatcher &sigc&
+waitmatch sigc '^Ready to catch signals; the magic word is (.*)$' MAGIC_WORD
+printargs ${MAGIC_WORD}
+cmp stdout args.want
+
 interrupt
 wait
-stdout 'caught interrupt'
+# Make sure the entire stdout still contains what waitmatch read.
+cmp stdout stdout.want
+
+-- args.want --
+["printargs" "Huzzah!"]
+-- stdout.want --
+Ready to catch signals; the magic word is Huzzah!
+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 &sigc&
+waitmatch sigc '^Ready to catch signals'

testscript/testscript.go

@@ -418,10 +418,27 @@ 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
+	stdoutReader io.ReadCloser    // a stdout pipe for waitmatch to read lines
+	stdoutBuffer *strings.Builder // all reads from stdoutReader are buffered here
+	test         io.ReadCloser
+	waitc        <-chan struct{}
+	neg          bool // if true, cmd should fail
+}
+
+func (b *backgroundCmd) wait() (stdout, stderr string) {
+	// Consume the rest of stdoutReader to fill stdoutBuffer.
+	io.Copy(io.Discard, b.stdoutReader)
+	b.stdoutReader.Close()
+	<-b.waitc
+	stdout = b.stdoutBuffer.String()
+	stderr = b.cmd.Stderr.(*strings.Builder).String()
+	return stdout, stderr
+}
+
+func (b *backgroundCmd) stderr() string {
+	return b.cmd.Stderr.(*strings.Builder).String()
 }
 
 func writeFile(name string, data []byte, perm fs.FileMode, excl bool) error {
@@ -575,7 +592,7 @@ func (ts *TestScript) run() {
 			ts.waitBackground(false)
 		} else {
 			for _, bg := range ts.background {
-				<-bg.wait
+				bg.wait()
 			}
 			ts.background = nil
 		}
@@ -1027,22 +1044,26 @@ 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) {
+func (ts *TestScript) execBackground(command string, args ...string) (*exec.Cmd, io.ReadCloser, error) {
 	if ts.ttyin != "" {
-		return nil, errors.New("ttyin is not supported by background commands")
+		return nil, nil, errors.New("ttyin is not supported by background commands")
 	}
 	cmd, err := ts.buildExecCmd(command, args...)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 	cmd.Dir = ts.cd
 	cmd.Env = append(ts.env, "PWD="+ts.cd)
-	var stdoutBuf, stderrBuf strings.Builder
+	var stderrBuf strings.Builder
 	cmd.Stdin = strings.NewReader(ts.stdin)
-	cmd.Stdout = &stdoutBuf
+	stdoutr, stdoutw, err := os.Pipe()
+	if err != nil {
+		return nil, nil, err
+	}
+	cmd.Stdout = stdoutw
 	cmd.Stderr = &stderrBuf
 	ts.stdin = ""
-	return cmd, cmd.Start()
+	return cmd, stdoutr, cmd.Start()
 }
 
 func (ts *TestScript) buildExecCmd(command string, args ...string) (*exec.Cmd, error) {
@@ -1143,6 +1164,9 @@ func waitOrStop(ctx context.Context, cmd *exec.Cmd, killDelay time.Duration) err
 	}()
 
 	waitErr := cmd.Wait()
+	if f, ok := cmd.Stdout.(*os.File); ok {
+		f.Close()
+	}
 	if interruptErr := <-errc; interruptErr != nil {
 		return interruptErr
 	}

testscript/testscript_test.go

@@ -46,12 +46,7 @@ func signalCatcher() int {
 	// Note: won't work under Windows.
 	c := make(chan os.Signal, 1)
 	signal.Notify(c, os.Interrupt)
-	// Create a file so that the test can know that
-	// we will catch the signal.
-	if err := os.WriteFile("catchsignal", nil, 0o666); err != nil {
-		fmt.Println(err)
-		return 1
-	}
+	fmt.Println("Ready to catch signals; the magic word is Huzzah!")
 	<-c
 	fmt.Println("caught interrupt")
 	return 0
@@ -180,7 +175,6 @@ func TestScripts(t *testing.T) {
 			"setSpecialVal":    setSpecialVal,
 			"ensureSpecialVal": ensureSpecialVal,
 			"interrupt":        interrupt,
-			"waitfile":         waitFile,
 			"testdefer": func(ts *TestScript, neg bool, args []string) {
 				testDeferCount++
 				n := testDeferCount
@@ -462,27 +456,6 @@ func interrupt(ts *TestScript, neg bool, args []string) {
 	bg[0].Process.Signal(os.Interrupt)
 }
 
-func waitFile(ts *TestScript, neg bool, args []string) {
-	if neg {
-		ts.Fatalf("waitfile does not support neg")
-	}
-	if len(args) != 1 {
-		ts.Fatalf("usage: waitfile file")
-	}
-	path := ts.MkAbs(args[0])
-	for i := 0; i < 100; i++ {
-		_, err := os.Stat(path)
-		if err == nil {
-			return
-		}
-		if !os.IsNotExist(err) {
-			ts.Fatalf("unexpected stat error: %v", err)
-		}
-		time.Sleep(10 * time.Millisecond)
-	}
-	ts.Fatalf("timed out waiting for %q to be created", path)
-}
-
 type fakeT struct {
 	log     strings.Builder
 	verbose bool