testscript: add "pty -stdin" option

testscript/cmd.go

@@ -392,6 +392,9 @@ func (ts *TestScript) cmdStdin(neg bool, args []string) {
 	if len(args) != 1 {
 		ts.Fatalf("usage: stdin filename")
 	}
+	if ts.stdinPty {
+		ts.Fatalf("conflicting use of 'stdin' and 'pty -stdin'")
+	}
 	ts.stdin = ts.ReadFile(args[0])
 }
 
@@ -418,10 +421,21 @@ func (ts *TestScript) cmdPty(neg bool, args []string) {
 	if neg {
 		ts.Fatalf("unsupported: ! pty")
 	}
-	if len(args) != 1 {
-		ts.Fatalf("usage: pty filename")
+	switch len(args) {
+	case 1:
+		ts.ptyin = ts.ReadFile(args[0])
+	case 2:
+		if args[0] != "-stdin" {
+			ts.Fatalf("usage: pty [-stdin] filename")
+		}
+		if ts.stdin != "" {
+			ts.Fatalf("conflicting use of 'stdin' and 'pty -stdin'")
+		}
+		ts.stdinPty = true
+		ts.ptyin = ts.ReadFile(args[1])
+	default:
+		ts.Fatalf("usage: pty [-stdin] filename")
 	}
-	ts.ptyin = ts.ReadFile(args[0])
 	if ts.ptyin == "" {
 		ts.Fatalf("pty input file is empty")
 	}

testscript/doc.go

@@ -198,10 +198,11 @@ The predefined commands are:
   Apply the grep command (see above) to the standard output
   from the most recent exec or wait command.
 
-- pty file
+- pty [-stdin] file
   Attach the next exec command to a controlling pseudo-terminal, and use the
-  contents of the given file as the raw terminal input.
-  Note that this does not attach the terminal to standard input/output/error.
+  contents of the given file as the raw terminal input. If -stdin is specified,
+  also attach the terminal to standard input.
+  Note that this does not attach the terminal to standard output/error.
 
 - [!] ptyout [-count=N] pattern
   Apply the grep command (see above) to the raw controlling terminal output

testscript/testscript.go

@@ -297,6 +297,7 @@ type TestScript struct {
 	stdout        string                      // standard output from last 'go' command; for 'stdout' command
 	stderr        string                      // standard error from last 'go' command; for 'stderr' command
 	ptyin         string                      // terminal input; set by 'pty' command
+	stdinPty      bool                        // connect pty to standard input; set by 'pty -stdin' command
 	ptyout        string                      // terminal output; for 'ptyout' command
 	stopped       bool                        // test wants to stop early
 	start         time.Time                   // time phase started
@@ -715,11 +716,15 @@ func (ts *TestScript) exec(command string, args ...string) (stdout, stderr strin
 		if err := setCtty(cmd, tty); err != nil {
 			return "", "", err
 		}
+		if ts.stdinPty {
+			cmd.Stdin = tty
+		}
 	}
 	if err = cmd.Start(); err == nil {
 		err = ctxWait(ts.ctxt, cmd)
 	}
 	ts.stdin = ""
+	ts.stdinPty = false
 	return stdoutBuf.String(), stderrBuf.String(), err
 }