diff --git a/agent/agent_test.go b/agent/agent_test.go index 2a8a2224ab..db32354179 100644 --- a/agent/agent_test.go +++ b/agent/agent_test.go @@ -1,6 +1,7 @@ package agent_test import ( + "bufio" "context" "encoding/json" "fmt" @@ -204,6 +205,11 @@ func TestAgent(t *testing.T) { id := uuid.NewString() netConn, err := conn.ReconnectingPTY(id, 100, 100) require.NoError(t, err) + bufRead := bufio.NewReader(netConn) + + // Brief pause to reduce the likelihood that we send keystrokes while + // the shell is simultaneously sending a prompt. + time.Sleep(100 * time.Millisecond) data, err := json.Marshal(agent.ReconnectingPTYRequest{ Data: "echo test\r\n", @@ -212,28 +218,35 @@ func TestAgent(t *testing.T) { _, err = netConn.Write(data) require.NoError(t, err) - findEcho := func() { + expectLine := func(matcher func(string) bool) { for { - read, err := netConn.Read(data) + line, err := bufRead.ReadString('\n') require.NoError(t, err) - if strings.Contains(string(data[:read]), "test") { + if matcher(line) { break } } } + matchEchoCommand := func(line string) bool { + return strings.Contains(line, "echo test") + } + matchEchoOutput := func(line string) bool { + return strings.Contains(line, "test") && !strings.Contains(line, "echo") + } // Once for typing the command... - findEcho() + expectLine(matchEchoCommand) // And another time for the actual output. - findEcho() + expectLine(matchEchoOutput) _ = netConn.Close() netConn, err = conn.ReconnectingPTY(id, 100, 100) require.NoError(t, err) + bufRead = bufio.NewReader(netConn) // Same output again! - findEcho() - findEcho() + expectLine(matchEchoCommand) + expectLine(matchEchoOutput) }) t.Run("Dial", func(t *testing.T) { diff --git a/coderd/workspaceagents_test.go b/coderd/workspaceagents_test.go index 09d3af4500..15bf01d065 100644 --- a/coderd/workspaceagents_test.go +++ b/coderd/workspaceagents_test.go @@ -1,11 +1,13 @@ package coderd_test import ( + "bufio" "context" "encoding/json" "runtime" "strings" "testing" + "time" "github.com/google/uuid" "github.com/pion/webrtc/v3" @@ -230,6 +232,11 @@ func TestWorkspaceAgentPTY(t *testing.T) { require.NoError(t, err) _, err = conn.Write(data) require.NoError(t, err) + bufRead := bufio.NewReader(conn) + + // Brief pause to reduce the likelihood that we send keystrokes while + // the shell is simultaneously sending a prompt. + time.Sleep(100 * time.Millisecond) data, err = json.Marshal(agent.ReconnectingPTYRequest{ Data: "echo test\r\n", @@ -238,16 +245,22 @@ func TestWorkspaceAgentPTY(t *testing.T) { _, err = conn.Write(data) require.NoError(t, err) - findEcho := func() { + expectLine := func(matcher func(string) bool) { for { - read, err := conn.Read(data) + line, err := bufRead.ReadString('\n') require.NoError(t, err) - if strings.Contains(string(data[:read]), "test") { - return + if matcher(line) { + break } } } + matchEchoCommand := func(line string) bool { + return strings.Contains(line, "echo test") + } + matchEchoOutput := func(line string) bool { + return strings.Contains(line, "test") && !strings.Contains(line, "echo") + } - findEcho() - findEcho() + expectLine(matchEchoCommand) + expectLine(matchEchoOutput) }