fix: Make TestAgent and TestWorkspaceAgentPTY less flaky (#1562)

dwahler · kylecarbs · commit c82c72473eae · 2022-06-10T17:01:39.000Z
diff --git 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
@@ -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)
 }
