fix: Resolve flake in TestPortForward (#4069)

cli/portforward_test.go

@@ -1,15 +1,12 @@
 package cli_test
 
 import (
-	"bytes"
 	"context"
 	"fmt"
 	"io"
 	"net"
-	"strings"
 	"sync"
 	"testing"
-	"time"
 
 	"github.com/google/uuid"
 	"github.com/pion/udp"
@@ -21,6 +18,7 @@ import (
 	"github.com/coder/coder/codersdk"
 	"github.com/coder/coder/provisioner/echo"
 	"github.com/coder/coder/provisionersdk/proto"
+	"github.com/coder/coder/pty/ptytest"
 	"github.com/coder/coder/testutil"
 )
 
@@ -35,15 +33,17 @@ func TestPortForward(t *testing.T) {
 
 		cmd, root := clitest.New(t, "port-forward", "blah")
 		clitest.SetupConfig(t, client, root)
-		buf := newThreadSafeBuffer()
-		cmd.SetOut(buf)
+		pty := ptytest.New(t)
+		cmd.SetIn(pty.Input())
+		cmd.SetOut(pty.Output())
+		cmd.SetErr(pty.Output())
 
 		err := cmd.Execute()
 		require.Error(t, err)
 		require.ErrorContains(t, err, "no port-forwards")
 
 		// Check that the help was printed.
-		require.Contains(t, buf.String(), "port-forward <workspace>")
+		pty.ExpectMatch("port-forward <workspace>")
 	})
 
 	cases := []struct {
@@ -135,15 +135,17 @@ func TestPortForward(t *testing.T) {
 				// the "local" listener.
 				cmd, root := clitest.New(t, "-v", "port-forward", workspace.Name, flag)
 				clitest.SetupConfig(t, client, root)
-				buf := newThreadSafeBuffer()
-				cmd.SetOut(buf)
+				pty := ptytest.New(t)
+				cmd.SetIn(pty.Input())
+				cmd.SetOut(pty.Output())
+				cmd.SetErr(pty.Output())
 				ctx, cancel := context.WithCancel(context.Background())
 				defer cancel()
 				errC := make(chan error)
 				go func() {
 					errC <- cmd.ExecuteContext(ctx)
 				}()
-				waitForPortForwardReady(t, buf)
+				pty.ExpectMatch("Ready!")
 
 				t.Parallel() // Port is reserved, enable parallel execution.
 
@@ -181,15 +183,17 @@ func TestPortForward(t *testing.T) {
 				// the "local" listeners.
 				cmd, root := clitest.New(t, "-v", "port-forward", workspace.Name, flag1, flag2)
 				clitest.SetupConfig(t, client, root)
-				buf := newThreadSafeBuffer()
-				cmd.SetOut(buf)
+				pty := ptytest.New(t)
+				cmd.SetIn(pty.Input())
+				cmd.SetOut(pty.Output())
+				cmd.SetErr(pty.Output())
 				ctx, cancel := context.WithCancel(context.Background())
 				defer cancel()
 				errC := make(chan error)
 				go func() {
 					errC <- cmd.ExecuteContext(ctx)
 				}()
-				waitForPortForwardReady(t, buf)
+				pty.ExpectMatch("Ready!")
 
 				t.Parallel() // Port is reserved, enable parallel execution.
 
@@ -236,15 +240,17 @@ func TestPortForward(t *testing.T) {
 		// the "local" listeners.
 		cmd, root := clitest.New(t, append([]string{"-v", "port-forward", workspace.Name}, flags...)...)
 		clitest.SetupConfig(t, client, root)
-		buf := newThreadSafeBuffer()
-		cmd.SetOut(buf)
+		pty := ptytest.New(t)
+		cmd.SetIn(pty.Input())
+		cmd.SetOut(pty.Output())
+		cmd.SetErr(pty.Output())
 		ctx, cancel := context.WithCancel(context.Background())
 		defer cancel()
 		errC := make(chan error)
 		go func() {
 			errC <- cmd.ExecuteContext(ctx)
 		}()
-		waitForPortForwardReady(t, buf)
+		pty.ExpectMatch("Ready!")
 
 		t.Parallel() // Port is reserved, enable parallel execution.
 
@@ -313,6 +319,10 @@ func runAgent(t *testing.T, client *codersdk.Client, userID uuid.UUID) ([]coders
 	// Start workspace agent in a goroutine
 	cmd, root := clitest.New(t, "agent", "--agent-token", agentToken, "--agent-url", client.URL.String())
 	clitest.SetupConfig(t, client, root)
+	pty := ptytest.New(t)
+	cmd.SetIn(pty.Input())
+	cmd.SetOut(pty.Output())
+	cmd.SetErr(pty.Output())
 	errC := make(chan error)
 	agentCtx, agentCancel := context.WithCancel(ctx)
 	t.Cleanup(func() {
@@ -404,61 +414,7 @@ func assertWritePayload(t *testing.T, w io.Writer, payload []byte) {
 	assert.Equal(t, len(payload), n, "payload length does not match")
 }
 
-func waitForPortForwardReady(t *testing.T, output *threadSafeBuffer) {
-	t.Helper()
-	for i := 0; i < 100; i++ {
-		time.Sleep(testutil.IntervalMedium)
-
-		data := output.String()
-		if strings.Contains(data, "Ready!") {
-			return
-		}
-	}
-
-	t.Fatal("port-forward command did not become ready in time")
-}
-
 type addr struct {
 	network string
 	addr    string
 }
-
-type threadSafeBuffer struct {
-	b   *bytes.Buffer
-	mut *sync.RWMutex
-}
-
-func newThreadSafeBuffer() *threadSafeBuffer {
-	return &threadSafeBuffer{
-		b:   bytes.NewBuffer(nil),
-		mut: new(sync.RWMutex),
-	}
-}
-
-var (
-	_ io.Reader = &threadSafeBuffer{}
-	_ io.Writer = &threadSafeBuffer{}
-)
-
-// Read implements io.Reader.
-func (b *threadSafeBuffer) Read(p []byte) (int, error) {
-	b.mut.RLock()
-	defer b.mut.RUnlock()
-
-	return b.b.Read(p)
-}
-
-// Write implements io.Writer.
-func (b *threadSafeBuffer) Write(p []byte) (int, error) {
-	b.mut.Lock()
-	defer b.mut.Unlock()
-
-	return b.b.Write(p)
-}
-
-func (b *threadSafeBuffer) String() string {
-	b.mut.RLock()
-	defer b.mut.RUnlock()
-
-	return b.b.String()
-}