testscript: support named background commands (#152)

This lets us wait for an individual background command rather
than all of them at once.

Author: rogpeppe

testscript/cmd.go

@@ -234,14 +234,20 @@ func (ts *TestScript) cmdEnv(neg bool, args []string) {
 	}
 }
 
+var backgroundSpecifier = regexp.MustCompile(`^&([a-zA-Z_0-9]+&)?$`)
+
 // exec runs the given command.
 func (ts *TestScript) cmdExec(neg bool, args []string) {
 	if len(args) < 1 || (len(args) == 1 && args[0] == "&") {
 		ts.Fatalf("usage: exec program [args...] [&]")
 	}
 
 	var err error
-	if len(args) > 0 && args[len(args)-1] == "&" {
+	if len(args) > 0 && backgroundSpecifier.MatchString(args[len(args)-1]) {
+		bgName := strings.TrimSuffix(strings.TrimPrefix(args[len(args)-1], "&"), "&")
+		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]...)
 		if err == nil {
@@ -250,7 +256,7 @@ func (ts *TestScript) cmdExec(neg bool, args []string) {
 				ctxWait(ts.ctxt, cmd)
 				close(wait)
 			}()
-			ts.background = append(ts.background, backgroundCmd{cmd, wait, neg})
+			ts.background = append(ts.background, backgroundCmd{bgName, cmd, wait, neg})
 		}
 		ts.stdout, ts.stderr = "", ""
 	} else {
@@ -449,16 +455,69 @@ func (ts *TestScript) cmdUNIX2DOS(neg bool, args []string) {
 
 // Tait waits for background commands to exit, setting stderr and stdout to their result.
 func (ts *TestScript) cmdWait(neg bool, args []string) {
+	if len(args) > 1 {
+		ts.Fatalf("usage: wait [name]")
+	}
 	if neg {
 		ts.Fatalf("unsupported: ! wait")
 	}
 	if len(args) > 0 {
-		ts.Fatalf("usage: wait")
+		ts.waitBackgroundOne(args[0])
+	} else {
+		ts.waitBackground(true)
+	}
+}
+
+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()
+	if ts.stdout != "" {
+		fmt.Fprintf(&ts.log, "[stdout]\n%s", ts.stdout)
+	}
+	if ts.stderr != "" {
+		fmt.Fprintf(&ts.log, "[stderr]\n%s", ts.stderr)
+	}
+	// Note: ignore bg.neg, which only takes effect on the non-specific
+	// wait command.
+	if bg.cmd.ProcessState.Success() {
+		if bg.neg {
+			ts.Fatalf("unexpected command success")
+		}
+	} else {
+		if ts.ctxt.Err() != nil {
+			ts.Fatalf("test timed out while running command")
+		} else if !bg.neg {
+			ts.Fatalf("unexpected command failure")
+		}
+	}
+	// Remove this process from the list of running background processes.
+	for i := range ts.background {
+		if bg == &ts.background[i] {
+			ts.background = append(ts.background[:i], ts.background[i+1:]...)
+			break
+		}
+	}
+}
+
+func (ts *TestScript) findBackground(bgName string) *backgroundCmd {
+	if bgName == "" {
+		return nil
+	}
+	for i := range ts.background {
+		bg := &ts.background[i]
+		if bg.name == bgName {
+			return bg
+		}
 	}
-	ts.waitBackground(true, neg)
+	return nil
 }
 
-func (ts *TestScript) waitBackground(checkStatus bool, neg bool) {
+func (ts *TestScript) waitBackground(checkStatus bool) {
 	var stdouts, stderrs []string
 	for _, bg := range ts.background {
 		<-bg.wait

testscript/doc.go

@@ -152,6 +152,10 @@ The predefined commands are:
   test. At the end of the test, any remaining background processes are
   terminated using os.Interrupt (if supported) or os.Kill.
 
+  If the last token is '&word&` (where "word" is alphanumeric), the
+  command runs in the background but has a name, and can be waited
+  for specifically by passing the word to 'wait'.
+
   Standard input can be provided using the stdin command; this will be
   cleared after exec has been called.
 
@@ -197,13 +201,15 @@ The predefined commands are:
 - symlink file -> target
   Create file as a symlink to target. The -> (like in ls -l output) is required.
 
-- wait
+- wait [command]
   Wait for all 'exec' and 'go' commands started in the background (with the '&'
   token) to exit, and display success or failure status for them.
   After a call to wait, the 'stderr' and 'stdout' commands will apply to the
   concatenation of the corresponding streams of the background commands,
   in the order in which those commands were started.
 
+  If an argument is specified, it waits for just that command.
+
 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/wait.txt

@@ -20,6 +20,21 @@ exec echo bar &
 wait
 stdout 'foo\nbar'
 
+exec echo bg1 &b1&
+exec echo bg2 &b2&
+exec echo bg3 &b3&
+exec echo bg4 &b4&
+
+wait b3
+stdout bg3
+wait b2
+stdout bg2
+wait
+stdout 'bg1\nbg4'
+
+# We should be able to start several background processes and wait for them
+# individually.
+
 # The end of the test should interrupt or kill any remaining background
 # programs.
 [!exec:sleep] skip

testscript/testscript.go

@@ -286,6 +286,7 @@ type TestScript struct {
 }
 
 type backgroundCmd struct {
+	name string
 	cmd  *exec.Cmd
 	wait <-chan struct{}
 	neg  bool // if true, cmd should fail
@@ -396,7 +397,7 @@ func (ts *TestScript) run() {
 		if ts.t.Verbose() || hasFailed(ts.t) {
 			// In verbose mode or on test failure, we want to see what happened in the background
 			// processes too.
-			ts.waitBackground(false, false)
+			ts.waitBackground(false)
 		} else {
 			for _, bg := range ts.background {
 				<-bg.wait