refactor: PTY & SSH (#7100)

* Add ssh tests for longoutput, orphan Signed-off-by: Spike Curtis <spike@coder.com> * PTY/SSH tests & improvements Signed-off-by: Spike Curtis <spike@coder.com> * Fix some tests Signed-off-by: Spike Curtis <spike@coder.com> * Fix linting Signed-off-by: Spike Curtis <spike@coder.com> * fmt Signed-off-by: Spike Curtis <spike@coder.com> * Fix windows test Signed-off-by: Spike Curtis <spike@coder.com> * Windows copy test Signed-off-by: Spike Curtis <spike@coder.com> * WIP Windows pty handling Signed-off-by: Spike Curtis <spike@coder.com> * Fix truncation tests Signed-off-by: Spike Curtis <spike@coder.com> * Appease linter/fmt Signed-off-by: Spike Curtis <spike@coder.com> * Fix typo Signed-off-by: Spike Curtis <spike@coder.com> * Rework truncation test to not assume OS buffers Signed-off-by: Spike Curtis <spike@coder.com> * Disable orphan test on Windows --- uses sh Signed-off-by: Spike Curtis <spike@coder.com> * agent_test running SSH in pty use ptytest.Start Signed-off-by: Spike Curtis <spike@coder.com> * More detail about closing pseudoconsole on windows Signed-off-by: Spike Curtis <spike@coder.com> * Code review fixes Signed-off-by: Spike Curtis <spike@coder.com> * Rearrange ptytest method order Signed-off-by: Spike Curtis <spike@coder.com> * Protect pty.Resize on windows from races Signed-off-by: Spike Curtis <spike@coder.com> * Fix windows bugs Signed-off-by: Spike Curtis <spike@coder.com> * PTY doesn't extend PTYCmd Signed-off-by: Spike Curtis <spike@coder.com> * Fix windows types Signed-off-by: Spike Curtis <spike@coder.com> --------- Signed-off-by: Spike Curtis <spike@coder.com>
2026-09-24 15:04:27 +08:00 · 2023-04-24 14:53:57 +04:00
parent c000f2ec28
commit daee91c6dc
16 changed files with 803 additions and 288 deletions
@@ -30,12 +30,21 @@ func New(t *testing.T, opts ...pty.Option) *PTY {
 	ptty, err := pty.New(opts...)
 	require.NoError(t, err)

-	return create(t, ptty, "cmd")
+	e := newExpecter(t, ptty.Output(), "cmd")
+	r := &PTY{
+		outExpecter: e,
+		PTY:         ptty,
+	}
+	// Ensure pty is cleaned up at the end of test.
+	t.Cleanup(func() {
+		_ = r.Close()
+	})
+	return r
 }

-// Start starts a new process asynchronously and returns a PTY and Process.
-// It kills the process upon cleanup.
-func Start(t *testing.T, cmd *exec.Cmd, opts ...pty.StartOption) (*PTY, pty.Process) {
+// Start starts a new process asynchronously and returns a PTYCmd and Process.
+// It kills the process and PTYCmd upon cleanup
+func Start(t *testing.T, cmd *exec.Cmd, opts ...pty.StartOption) (*PTYCmd, pty.Process) {
 	t.Helper()

 	ptty, ps, err := pty.Start(cmd, opts...)
@@ -44,10 +53,19 @@ func Start(t *testing.T, cmd *exec.Cmd, opts ...pty.StartOption) (*PTY, pty.Proc
 		_ = ps.Kill()
 		_ = ps.Wait()
 	})
-	return create(t, ptty, cmd.Args[0]), ps
+	ex := newExpecter(t, ptty.OutputReader(), cmd.Args[0])
+
+	r := &PTYCmd{
+		outExpecter: ex,
+		PTYCmd:      ptty,
+	}
+	t.Cleanup(func() {
+		_ = r.Close()
+	})
+	return r, ps
 }

-func create(t *testing.T, ptty pty.PTY, name string) *PTY {
+func newExpecter(t *testing.T, r io.Reader, name string) outExpecter {
 	// Use pipe for logging.
 	logDone := make(chan struct{})
 	logr, logw := io.Pipe()
@@ -57,37 +75,30 @@ func create(t *testing.T, ptty pty.PTY, name string) *PTY {
 	out := newStdbuf()
 	w := io.MultiWriter(logw, out)

-	tpty := &PTY{
+	ex := outExpecter{
 		t:    t,
-		PTY:  ptty,
 		out:  out,
 		name: name,

 		runeReader: bufio.NewReaderSize(out, utf8.UTFMax),
 	}
-	// Ensure pty is cleaned up at the end of test.
-	t.Cleanup(func() {
-		_ = tpty.Close()
-	})

 	logClose := func(name string, c io.Closer) {
-		tpty.logf("closing %s", name)
+		ex.logf("closing %s", name)
 		err := c.Close()
-		tpty.logf("closed %s: %v", name, err)
+		ex.logf("closed %s: %v", name, err)
 	}
-	// Set the actual close function for the tpty.
-	tpty.close = func(reason string) error {
+	// Set the actual close function for the outExpecter.
+	ex.close = func(reason string) error {
 		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
 		defer cancel()

-		tpty.logf("closing tpty: %s", reason)
+		ex.logf("closing expecter: %s", reason)

-		// Close pty only so that the copy goroutine can consume the
-		// remainder of it's buffer and then exit.
-		logClose("pty", ptty)
+		// Caller needs to have closed the PTY so that copying can complete
 		select {
 		case <-ctx.Done():
-			tpty.fatalf("close", "copy did not close in time")
+			ex.fatalf("close", "copy did not close in time")
 		case <-copyDone:
 		}

@@ -95,22 +106,22 @@ func create(t *testing.T, ptty pty.PTY, name string) *PTY {
 		logClose("logr", logr)
 		select {
 		case <-ctx.Done():
-			tpty.fatalf("close", "log pipe did not close in time")
+			ex.fatalf("close", "log pipe did not close in time")
 		case <-logDone:
 		}

-		tpty.logf("closed tpty")
+		ex.logf("closed expecter")

 		return nil
 	}

 	go func() {
 		defer close(copyDone)
-		_, err := io.Copy(w, ptty.Output())
-		tpty.logf("copy done: %v", err)
-		tpty.logf("closing out")
+		_, err := io.Copy(w, r)
+		ex.logf("copy done: %v", err)
+		ex.logf("closing out")
 		err = out.closeErr(err)
-		tpty.logf("closed out: %v", err)
+		ex.logf("closed out: %v", err)
 	}()

 	// Log all output as part of test for easier debugging on errors.
@@ -118,15 +129,14 @@ func create(t *testing.T, ptty pty.PTY, name string) *PTY {
 		defer close(logDone)
 		s := bufio.NewScanner(logr)
 		for s.Scan() {
-			tpty.logf("%q", stripansi.Strip(s.Text()))
+			ex.logf("%q", stripansi.Strip(s.Text()))
 		}
 	}()

-	return tpty
+	return ex
 }

-type PTY struct {
-	pty.PTY
+type outExpecter struct {
 	t     *testing.T
 	close func(reason string) error
 	out   *stdbuf
@@ -135,38 +145,23 @@ type PTY struct {
 	runeReader *bufio.Reader
 }

-func (p *PTY) Close() error {
-	p.t.Helper()
-
-	return p.close("close")
-}
-
-func (p *PTY) Attach(inv *clibase.Invocation) *PTY {
-	p.t.Helper()
-
-	inv.Stdout = p.Output()
-	inv.Stderr = p.Output()
-	inv.Stdin = p.Input()
-	return p
-}
-
-func (p *PTY) ExpectMatch(str string) string {
-	p.t.Helper()
+func (e *outExpecter) ExpectMatch(str string) string {
+	e.t.Helper()

 	timeout, cancel := context.WithTimeout(context.Background(), testutil.WaitMedium)
 	defer cancel()

-	return p.ExpectMatchContext(timeout, str)
+	return e.ExpectMatchContext(timeout, str)
 }

 // TODO(mafredri): Rename this to ExpectMatch when refactoring.
-func (p *PTY) ExpectMatchContext(ctx context.Context, str string) string {
-	p.t.Helper()
+func (e *outExpecter) ExpectMatchContext(ctx context.Context, str string) string {
+	e.t.Helper()

 	var buffer bytes.Buffer
-	err := p.doMatchWithDeadline(ctx, "ExpectMatchContext", func() error {
+	err := e.doMatchWithDeadline(ctx, "ExpectMatchContext", func() error {
 		for {
-			r, _, err := p.runeReader.ReadRune()
+			r, _, err := e.runeReader.ReadRune()
 			if err != nil {
 				return err
 			}
@@ -180,54 +175,54 @@ func (p *PTY) ExpectMatchContext(ctx context.Context, str string) string {
 		}
 	})
 	if err != nil {
-		p.fatalf("read error", "%v (wanted %q; got %q)", err, str, buffer.String())
+		e.fatalf("read error", "%v (wanted %q; got %q)", err, str, buffer.String())
 		return ""
 	}
-	p.logf("matched %q = %q", str, stripansi.Strip(buffer.String()))
+	e.logf("matched %q = %q", str, stripansi.Strip(buffer.String()))
 	return buffer.String()
 }

-func (p *PTY) Peek(ctx context.Context, n int) []byte {
-	p.t.Helper()
+func (e *outExpecter) Peek(ctx context.Context, n int) []byte {
+	e.t.Helper()

 	var out []byte
-	err := p.doMatchWithDeadline(ctx, "Peek", func() error {
+	err := e.doMatchWithDeadline(ctx, "Peek", func() error {
 		var err error
-		out, err = p.runeReader.Peek(n)
+		out, err = e.runeReader.Peek(n)
 		return err
 	})
 	if err != nil {
-		p.fatalf("read error", "%v (wanted %d bytes; got %d: %q)", err, n, len(out), out)
+		e.fatalf("read error", "%v (wanted %d bytes; got %d: %q)", err, n, len(out), out)
 		return nil
 	}
-	p.logf("peeked %d/%d bytes = %q", len(out), n, out)
+	e.logf("peeked %d/%d bytes = %q", len(out), n, out)
 	return slices.Clone(out)
 }

-func (p *PTY) ReadRune(ctx context.Context) rune {
-	p.t.Helper()
+func (e *outExpecter) ReadRune(ctx context.Context) rune {
+	e.t.Helper()

 	var r rune
-	err := p.doMatchWithDeadline(ctx, "ReadRune", func() error {
+	err := e.doMatchWithDeadline(ctx, "ReadRune", func() error {
 		var err error
-		r, _, err = p.runeReader.ReadRune()
+		r, _, err = e.runeReader.ReadRune()
 		return err
 	})
 	if err != nil {
-		p.fatalf("read error", "%v (wanted rune; got %q)", err, r)
+		e.fatalf("read error", "%v (wanted rune; got %q)", err, r)
 		return 0
 	}
-	p.logf("matched rune = %q", r)
+	e.logf("matched rune = %q", r)
 	return r
 }

-func (p *PTY) ReadLine(ctx context.Context) string {
-	p.t.Helper()
+func (e *outExpecter) ReadLine(ctx context.Context) string {
+	e.t.Helper()

 	var buffer bytes.Buffer
-	err := p.doMatchWithDeadline(ctx, "ReadLine", func() error {
+	err := e.doMatchWithDeadline(ctx, "ReadLine", func() error {
 		for {
-			r, _, err := p.runeReader.ReadRune()
+			r, _, err := e.runeReader.ReadRune()
 			if err != nil {
 				return err
 			}
@@ -240,14 +235,14 @@ func (p *PTY) ReadLine(ctx context.Context) string {

 				// Unicode code points can be up to 4 bytes, but the
 				// ones we're looking for are only 1 byte.
-				b, _ := p.runeReader.Peek(1)
+				b, _ := e.runeReader.Peek(1)
 				if len(b) == 0 {
 					return nil
 				}

 				r, _ = utf8.DecodeRune(b)
 				if r == '\n' {
-					_, _, err = p.runeReader.ReadRune()
+					_, _, err = e.runeReader.ReadRune()
 					if err != nil {
 						return err
 					}
@@ -263,21 +258,21 @@ func (p *PTY) ReadLine(ctx context.Context) string {
 		}
 	})
 	if err != nil {
-		p.fatalf("read error", "%v (wanted newline; got %q)", err, buffer.String())
+		e.fatalf("read error", "%v (wanted newline; got %q)", err, buffer.String())
 		return ""
 	}
-	p.logf("matched newline = %q", buffer.String())
+	e.logf("matched newline = %q", buffer.String())
 	return buffer.String()
 }

-func (p *PTY) doMatchWithDeadline(ctx context.Context, name string, fn func() error) error {
-	p.t.Helper()
+func (e *outExpecter) doMatchWithDeadline(ctx context.Context, name string, fn func() error) error {
+	e.t.Helper()

 	// A timeout is mandatory, caller can decide by passing a context
 	// that times out.
 	if _, ok := ctx.Deadline(); !ok {
 		timeout := testutil.WaitMedium
-		p.logf("%s ctx has no deadline, using %s", name, timeout)
+		e.logf("%s ctx has no deadline, using %s", name, timeout)
 		var cancel context.CancelFunc
 		//nolint:gocritic // Rule guard doesn't detect that we're using testutil.Wait*.
 		ctx, cancel = context.WithTimeout(ctx, timeout)
@@ -294,13 +289,55 @@ func (p *PTY) doMatchWithDeadline(ctx context.Context, name string, fn func() er
 		return err
 	case <-ctx.Done():
 		// Ensure goroutine is cleaned up before test exit.
-		_ = p.close("match deadline exceeded")
+		_ = e.close("match deadline exceeded")
 		<-match

 		return xerrors.Errorf("match deadline exceeded: %w", ctx.Err())
 	}
 }

+func (e *outExpecter) logf(format string, args ...interface{}) {
+	e.t.Helper()
+
+	// Match regular logger timestamp format, we seem to be logging in
+	// UTC in other places as well, so match here.
+	e.t.Logf("%s: %s: %s", time.Now().UTC().Format("2006-01-02 15:04:05.000"), e.name, fmt.Sprintf(format, args...))
+}
+
+func (e *outExpecter) fatalf(reason string, format string, args ...interface{}) {
+	e.t.Helper()
+
+	// Ensure the message is part of the normal log stream before
+	// failing the test.
+	e.logf("%s: %s", reason, fmt.Sprintf(format, args...))
+
+	require.FailNowf(e.t, reason, format, args...)
+}
+
+type PTY struct {
+	outExpecter
+	pty.PTY
+}
+
+func (p *PTY) Close() error {
+	p.t.Helper()
+	pErr := p.PTY.Close()
+	eErr := p.outExpecter.close("close")
+	if pErr != nil {
+		return pErr
+	}
+	return eErr
+}
+
+func (p *PTY) Attach(inv *clibase.Invocation) *PTY {
+	p.t.Helper()
+
+	inv.Stdout = p.Output()
+	inv.Stderr = p.Output()
+	inv.Stdin = p.Input()
+	return p
+}
+
 func (p *PTY) Write(r rune) {
 	p.t.Helper()

@@ -321,22 +358,19 @@ func (p *PTY) WriteLine(str string) {
 	require.NoError(p.t, err, "write line failed")
 }

-func (p *PTY) logf(format string, args ...interface{}) {
-	p.t.Helper()
-
-	// Match regular logger timestamp format, we seem to be logging in
-	// UTC in other places as well, so match here.
-	p.t.Logf("%s: %s: %s", time.Now().UTC().Format("2006-01-02 15:04:05.000"), p.name, fmt.Sprintf(format, args...))
+type PTYCmd struct {
+	outExpecter
+	pty.PTYCmd
 }

-func (p *PTY) fatalf(reason string, format string, args ...interface{}) {
+func (p *PTYCmd) Close() error {
 	p.t.Helper()
-
-	// Ensure the message is part of the normal log stream before
-	// failing the test.
-	p.logf("%s: %s", reason, fmt.Sprintf(format, args...))
-
-	require.FailNowf(p.t, reason, format, args...)
+	pErr := p.PTYCmd.Close()
+	eErr := p.outExpecter.close("close")
+	if pErr != nil {
+		return pErr
+	}
+	return eErr
 }

 // stdbuf is like a buffered stdout, it buffers writes until read.