fix: use terminal emulator that keeps state in ReconnectingPTY tests (#9765)

* Add more pty diagnostics for terminal parsing

Signed-off-by: Spike Curtis <spike@coder.com>

* print escaped strings

Signed-off-by: Spike Curtis <spike@coder.com>

* Only log on failure - heisenbug?

Signed-off-by: Spike Curtis <spike@coder.com>

* use the terminal across matches to keep cursor & contents state

Signed-off-by: Spike Curtis <spike@coder.com>

* Only log bytes if we're not expecting EOF

Signed-off-by: Spike Curtis <spike@coder.com>

---------

Signed-off-by: Spike Curtis <spike@coder.com>

agent/agent_test.go

@@ -1669,13 +1669,15 @@ func TestAgent_ReconnectingPTY(t *testing.T) {
 			}
 
 			// Once for typing the command...
-			require.NoError(t, testutil.ReadUntil(ctx, t, netConn1, matchEchoCommand), "find echo command")
+			tr1 := testutil.NewTerminalReader(t, netConn1)
+			require.NoError(t, tr1.ReadUntil(ctx, matchEchoCommand), "find echo command")
 			// And another time for the actual output.
-			require.NoError(t, testutil.ReadUntil(ctx, t, netConn1, matchEchoOutput), "find echo output")
+			require.NoError(t, tr1.ReadUntil(ctx, matchEchoOutput), "find echo output")
 
 			// Same for the other connection.
-			require.NoError(t, testutil.ReadUntil(ctx, t, netConn2, matchEchoCommand), "find echo command")
+			tr2 := testutil.NewTerminalReader(t, netConn2)
-			require.NoError(t, testutil.ReadUntil(ctx, t, netConn2, matchEchoOutput), "find echo output")
+			require.NoError(t, tr2.ReadUntil(ctx, matchEchoCommand), "find echo command")
+			require.NoError(t, tr2.ReadUntil(ctx, matchEchoOutput), "find echo output")
 
 			_ = netConn1.Close()
 			_ = netConn2.Close()
@@ -1684,8 +1686,9 @@ func TestAgent_ReconnectingPTY(t *testing.T) {
 			defer netConn3.Close()
 
 			// Same output again!
-			require.NoError(t, testutil.ReadUntil(ctx, t, netConn3, matchEchoCommand), "find echo command")
+			tr3 := testutil.NewTerminalReader(t, netConn3)
-			require.NoError(t, testutil.ReadUntil(ctx, t, netConn3, matchEchoOutput), "find echo output")
+			require.NoError(t, tr3.ReadUntil(ctx, matchEchoCommand), "find echo command")
+			require.NoError(t, tr3.ReadUntil(ctx, matchEchoOutput), "find echo output")
 
 			// Exit should cause the connection to close.
 			data, err = json.Marshal(codersdk.ReconnectingPTYRequest{
@@ -1696,19 +1699,20 @@ func TestAgent_ReconnectingPTY(t *testing.T) {
 			require.NoError(t, err)
 
 			// Once for the input and again for the output.
-			require.NoError(t, testutil.ReadUntil(ctx, t, netConn3, matchExitCommand), "find exit command")
+			require.NoError(t, tr3.ReadUntil(ctx, matchExitCommand), "find exit command")
-			require.NoError(t, testutil.ReadUntil(ctx, t, netConn3, matchExitOutput), "find exit output")
+			require.NoError(t, tr3.ReadUntil(ctx, matchExitOutput), "find exit output")
 
 			// Wait for the connection to close.
-			require.ErrorIs(t, testutil.ReadUntil(ctx, t, netConn3, nil), io.EOF)
+			require.ErrorIs(t, tr3.ReadUntil(ctx, nil), io.EOF)
 
 			// Try a non-shell command.  It should output then immediately exit.
 			netConn4, err := conn.ReconnectingPTY(ctx, uuid.New(), 80, 80, "echo test")
 			require.NoError(t, err)
 			defer netConn4.Close()
 
-			require.NoError(t, testutil.ReadUntil(ctx, t, netConn4, matchEchoOutput), "find echo output")
+			tr4 := testutil.NewTerminalReader(t, netConn4)
-			require.ErrorIs(t, testutil.ReadUntil(ctx, t, netConn3, nil), io.EOF)
+			require.NoError(t, tr4.ReadUntil(ctx, matchEchoOutput), "find echo output")
+			require.ErrorIs(t, tr4.ReadUntil(ctx, nil), io.EOF)
 		})
 	}
 }

coderd/workspaceapps/apptest/apptest.go

@@ -1428,8 +1428,9 @@ func testReconnectingPTY(ctx context.Context, t *testing.T, client *codersdk.Cli
 	_, err = conn.Write(data)
 	require.NoError(t, err)
 
-	require.NoError(t, testutil.ReadUntil(ctx, t, conn, matchEchoCommand), "find echo command")
+	tr := testutil.NewTerminalReader(t, conn)
-	require.NoError(t, testutil.ReadUntil(ctx, t, conn, matchEchoOutput), "find echo output")
+	require.NoError(t, tr.ReadUntil(ctx, matchEchoCommand), "find echo command")
+	require.NoError(t, tr.ReadUntil(ctx, matchEchoOutput), "find echo output")
 
 	// Exit should cause the connection to close.
 	data, err = json.Marshal(codersdk.ReconnectingPTYRequest{
@@ -1440,9 +1441,9 @@ func testReconnectingPTY(ctx context.Context, t *testing.T, client *codersdk.Cli
 	require.NoError(t, err)
 
 	// Once for the input and again for the output.
-	require.NoError(t, testutil.ReadUntil(ctx, t, conn, matchExitCommand), "find exit command")
+	require.NoError(t, tr.ReadUntil(ctx, matchExitCommand), "find exit command")
-	require.NoError(t, testutil.ReadUntil(ctx, t, conn, matchExitOutput), "find exit output")
+	require.NoError(t, tr.ReadUntil(ctx, matchExitOutput), "find exit output")
 
 	// Ensure the connection closes.
-	require.ErrorIs(t, testutil.ReadUntil(ctx, t, conn, nil), io.EOF)
+	require.ErrorIs(t, tr.ReadUntil(ctx, nil), io.EOF)
 }

pty/start_test.go

@@ -67,11 +67,11 @@ func Test_Start_truncation(t *testing.T) {
 	readDone := make(chan struct{})
 	go func() {
 		defer close(readDone)
-		// avoid buffered IO so that we can precisely control how many bytes to read.
+		terminalReader := testutil.NewTerminalReader(t, pc.OutputReader())
 		n := 1
 		for n <= countEnd {
 			want := fmt.Sprintf("%d", n)
-			err := testutil.ReadUntilString(ctx, t, want, pc.OutputReader())
+			err := terminalReader.ReadUntilString(ctx, want)
 			assert.NoError(t, err, "want: %s", want)
 			if err != nil {
 				return

testutil/pty.go

@@ -6,14 +6,33 @@ import (
 	"strings"
 	"testing"
 
+	"golang.org/x/xerrors"
+
 	"github.com/hinshun/vt10x"
 )
 
+// TerminalReader emulates a terminal and allows matching output.  It's important in cases where we
+// can get control sequences to parse them correctly, and keep the state of the terminal across the
+// lifespan of the PTY, since some control sequences are relative to the current cursor position.
+type TerminalReader struct {
+	t    *testing.T
+	r    io.Reader
+	term vt10x.Terminal
+}
+
+func NewTerminalReader(t *testing.T, r io.Reader) *TerminalReader {
+	return &TerminalReader{
+		t:    t,
+		r:    r,
+		term: vt10x.New(vt10x.WithSize(80, 80)),
+	}
+}
+
 // ReadUntilString emulates a terminal and reads one byte at a time until we
 // either see the string we want, or the context expires.  The PTY must be sized
 // to 80x80 or there could be unexpected results.
-func ReadUntilString(ctx context.Context, t *testing.T, want string, r io.Reader) error {
+func (tr *TerminalReader) ReadUntilString(ctx context.Context, want string) error {
-	return ReadUntil(ctx, t, r, func(line string) bool {
+	return tr.ReadUntil(ctx, func(line string) bool {
 		return strings.TrimSpace(line) == want
 	})
 }
@@ -21,26 +40,37 @@ func ReadUntilString(ctx context.Context, t *testing.T, want string, r io.Reader
 // ReadUntil emulates a terminal and reads one byte at a time until the matcher
 // returns true or the context expires.  If the matcher is nil, read until EOF.
 // The PTY must be sized to 80x80 or there could be unexpected results.
-func ReadUntil(ctx context.Context, t *testing.T, r io.Reader, matcher func(line string) bool) error {
+func (tr *TerminalReader) ReadUntil(ctx context.Context, matcher func(line string) bool) (retErr error) {
-	// output can contain virtual terminal sequences, so we need to parse these
+	readBytes := make([]byte, 0)
-	// to correctly interpret getting what we want.
-	term := vt10x.New(vt10x.WithSize(80, 80))
 	readErrs := make(chan error, 1)
 	defer func() {
 		// Dump the terminal contents since they can be helpful for debugging, but
-		// skip empty lines since much of the terminal will usually be blank.
+		// trim empty lines since much of the terminal will usually be blank.
-		got := term.String()
+		got := tr.term.String()
 		lines := strings.Split(got, "\n")
-		for _, line := range lines {
+		for i := range lines {
-			if strings.TrimSpace(line) != "" {
+			if strings.TrimSpace(lines[i]) != "" {
-				t.Logf("got: %v", line)
+				lines = lines[i:]
+				break
 			}
 		}
+		for i := len(lines) - 1; i >= 0; i-- {
+			if strings.TrimSpace(lines[i]) != "" {
+				lines = lines[:i+1]
+				break
+			}
+		}
+		gotTrimmed := strings.Join(lines, "\n")
+		tr.t.Logf("Terminal contents:\n%s", gotTrimmed)
+		// EOF is expected when matcher == nil
+		if retErr != nil && !(xerrors.Is(retErr, io.EOF) && matcher == nil) {
+			tr.t.Logf("Bytes Read: %q", string(readBytes))
+		}
 	}()
 	for {
 		b := make([]byte, 1)
 		go func() {
-			_, err := r.Read(b)
+			_, err := tr.r.Read(b)
 			readErrs <- err
 		}()
 		select {
@@ -48,7 +78,8 @@ func ReadUntil(ctx context.Context, t *testing.T, r io.Reader, matcher func(line
 			if err != nil {
 				return err
 			}
-			_, err = term.Write(b)
+			readBytes = append(readBytes, b...)
+			_, err = tr.term.Write(b)
 			if err != nil {
 				return err
 			}
@@ -59,7 +90,7 @@ func ReadUntil(ctx context.Context, t *testing.T, r io.Reader, matcher func(line
 			// A nil matcher means to read until EOF.
 			continue
 		}
-		got := term.String()
+		got := tr.term.String()
 		lines := strings.Split(got, "\n")
 		for _, line := range lines {
 			if matcher(line) {