diff --git a/agent/agent_test.go b/agent/agent_test.go index 9fbe263faf..bf234b5cfb 100644 --- a/agent/agent_test.go +++ b/agent/agent_test.go @@ -54,10 +54,10 @@ import ( "github.com/coder/coder/v2/codersdk/agentsdk" "github.com/coder/coder/v2/codersdk/workspacesdk" "github.com/coder/coder/v2/cryptorand" - "github.com/coder/coder/v2/pty/ptytest" "github.com/coder/coder/v2/tailnet" "github.com/coder/coder/v2/tailnet/tailnettest" "github.com/coder/coder/v2/testutil" + "github.com/coder/coder/v2/testutil/expecter" "github.com/coder/quartz" ) @@ -721,6 +721,7 @@ func TestAgent_SessionTTYShell(t *testing.T) { t.Run(fmt.Sprintf("(%d)", port), func(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitMedium) + logger := testutil.Logger(t) session := setupSSHSessionOnPort(t, agentsdk.Manifest{}, codersdk.ServiceBannerConfig{}, nil, port) command := "sh" @@ -729,16 +730,14 @@ func TestAgent_SessionTTYShell(t *testing.T) { } err := session.RequestPty("xterm", 128, 128, ssh.TerminalModes{}) require.NoError(t, err) - ptty := ptytest.New(t) - session.Stdout = ptty.Output() - session.Stderr = ptty.Output() - session.Stdin = ptty.Input() + stdout := expecter.NewAttachedToSSHSession(t, session) + stdin := testutil.NewWriterAttachedToSSHSession(t, logger.Named("sshin"), session) err = session.Start(command) require.NoError(t, err) - _ = ptty.Peek(ctx, 1) // wait for the prompt - ptty.WriteLine("echo test") - ptty.ExpectMatch(ctx, "test") - ptty.WriteLine("exit") + _ = stdout.Peek(ctx, 1) // wait for the prompt + stdin.WriteLine("echo test") + stdout.ExpectMatch(ctx, "test") + stdin.WriteLine("exit") err = session.Wait() require.NoError(t, err) }) @@ -751,10 +750,6 @@ func TestAgent_SessionTTYExitCode(t *testing.T) { command := "areallynotrealcommand" err := session.RequestPty("xterm", 128, 128, ssh.TerminalModes{}) require.NoError(t, err) - ptty := ptytest.New(t) - session.Stdout = ptty.Output() - session.Stderr = ptty.Output() - session.Stdin = ptty.Input() err = session.Start(command) require.NoError(t, err) err = session.Wait() @@ -1035,9 +1030,8 @@ func TestAgent_Session_TTY_QuietLogin(t *testing.T) { require.NoError(t, err) stdout := testutil.NewWaitBuffer() - ptty := ptytest.New(t) + session.Stdout = stdout - session.Stderr = ptty.Output() stdin, err := session.StdinPipe() require.NoError(t, err) require.NoError(t, session.Shell()) @@ -1076,8 +1070,6 @@ func TestAgent_Session_TTY_FastCommandHasOutput(t *testing.T) { require.NoError(t, err) defer sshClient.Close() - ptty := ptytest.New(t) - var stdout bytes.Buffer // NOTE(mafredri): Increase iterations to increase chance of failure, // assuming bug is present. Limiting GOMAXPROCS further @@ -1097,8 +1089,6 @@ func TestAgent_Session_TTY_FastCommandHasOutput(t *testing.T) { require.NoError(t, err) session.Stdout = &stdout - session.Stderr = ptty.Output() - session.Stdin = ptty.Input() err = session.Start("echo wazzup") require.NoError(t, err) @@ -1126,8 +1116,6 @@ func TestAgent_Session_TTY_HugeOutputIsNotLost(t *testing.T) { require.NoError(t, err) defer sshClient.Close() - ptty := ptytest.New(t) - var stdout bytes.Buffer // NOTE(mafredri): Increase iterations to increase chance of failure, // assuming bug is present. @@ -1146,8 +1134,6 @@ func TestAgent_Session_TTY_HugeOutputIsNotLost(t *testing.T) { require.NoError(t, err) session.Stdout = &stdout - session.Stderr = ptty.Output() - session.Stdin = ptty.Input() want := strings.Repeat("wazzup", 1024+1) // ~6KB, +1 because 1024 is a common buffer size. err = session.Start("echo " + want) require.NoError(t, err) @@ -4220,19 +4206,18 @@ func assertWritePayload(t testing.TB, w io.Writer, payload []byte) { func testSessionOutput(t *testing.T, session *ssh.Session, expected, unexpected []string, expectedRe *regexp.Regexp) { t.Helper() + logger := testutil.Logger(t) err := session.RequestPty("xterm", 128, 128, ssh.TerminalModes{}) require.NoError(t, err) - ptty := ptytest.New(t) var stdout bytes.Buffer session.Stdout = &stdout - session.Stderr = ptty.Output() - session.Stdin = ptty.Input() + stdin := testutil.NewWriterAttachedToSSHSession(t, logger.Named("sshin"), session) err = session.Shell() require.NoError(t, err) - ptty.WriteLine("exit 0") + stdin.WriteLine("exit 0") waitErr := make(chan error, 1) go func() { diff --git a/agent/agentssh/agentssh_test.go b/agent/agentssh/agentssh_test.go index fceed50abe..589b8c5510 100644 --- a/agent/agentssh/agentssh_test.go +++ b/agent/agentssh/agentssh_test.go @@ -28,8 +28,8 @@ import ( "cdr.dev/slog/v3/sloggers/slogtest" "github.com/coder/coder/v2/agent/agentexec" "github.com/coder/coder/v2/agent/agentssh" - "github.com/coder/coder/v2/pty/ptytest" "github.com/coder/coder/v2/testutil" + "github.com/coder/coder/v2/testutil/expecter" ) func TestMain(m *testing.M) { @@ -182,10 +182,8 @@ func TestNewServer_CloseActiveConnections(t *testing.T) { c := sshClient(t, ln.Addr().String()) sess, err := c.NewSession() assert.NoError(t, err) - pty := ptytest.New(t) - sess.Stdin = pty.Input() - sess.Stdout = pty.Output() - sess.Stderr = pty.Output() + stdout := expecter.NewAttachedToSSHSession(t, sess) + stdout.Rename(fmt.Sprintf("sess%d", i)) // Every other session will request a PTY. if i%2 == 0 { @@ -203,7 +201,7 @@ func TestNewServer_CloseActiveConnections(t *testing.T) { assert.NoError(t, err) // Allow the session to settle (i.e. reach echo). - pty.ExpectMatch(ctx, "started") + stdout.ExpectMatch(ctx, "started") // Sleep a bit to ensure the sleep has started. time.Sleep(testutil.IntervalMedium) @@ -353,18 +351,11 @@ func TestNewServer_Signal(t *testing.T) { c := sshClient(t, ln.Addr().String()) - pty := ptytest.New(t) - sess, err := c.NewSession() require.NoError(t, err) r, err := sess.StdoutPipe() require.NoError(t, err) - // Note, we request pty but don't use ptytest here because we can't - // easily test for no text before EOF. - sess.Stdin = pty.Input() - sess.Stderr = pty.Output() - err = sess.RequestPty("xterm", 80, 80, nil) require.NoError(t, err) diff --git a/testutil/expecter/expecter.go b/testutil/expecter/expecter.go index 858c82467a..d3f30c7a05 100644 --- a/testutil/expecter/expecter.go +++ b/testutil/expecter/expecter.go @@ -16,6 +16,7 @@ import ( "github.com/acarl005/stripansi" "github.com/stretchr/testify/require" "go.uber.org/atomic" + "golang.org/x/crypto/ssh" "golang.org/x/xerrors" "github.com/coder/coder/v2/testutil" @@ -149,6 +150,14 @@ func NewPiped(t *testing.T) (*Expecter, io.Writer) { return e, w } +func NewAttachedToSSHSession(t *testing.T, session *ssh.Session) *Expecter { + e, w := NewPiped(t) + e.Rename("sshout") + session.Stdout = w + session.Stderr = w + return e +} + type Expecter struct { t *testing.T out *stdbuf diff --git a/testutil/writer.go b/testutil/writer.go index 4def987e62..7fd288b5b0 100644 --- a/testutil/writer.go +++ b/testutil/writer.go @@ -5,6 +5,7 @@ import ( "testing" "github.com/stretchr/testify/assert" + "golang.org/x/crypto/ssh" "gvisor.dev/gvisor/pkg/context" "cdr.dev/slog/v3" @@ -32,6 +33,20 @@ func NewWriterAttachedToInvocation(t *testing.T, logger slog.Logger, invocation } } +func NewWriterAttachedToSSHSession(t *testing.T, l slog.Logger, session *ssh.Session) *Writer { + r, w := io.Pipe() + session.Stdin = r + // Close the pipe at the end of the test to ensure any goroutine in the Invocation that reads from stdin won't leak. + t.Cleanup(func() { + _ = w.Close() + }) + return &Writer{ + t: t, + w: w, + l: l, + } +} + func (w *Writer) Write(r rune) { w.t.Helper() _, err := w.w.Write([]byte{byte(r)})