test(agentcompat): add protocol clients

Co-authored-by: naiba/CloudCode <hi+cloudcode@nai.ba>

integration/agentcompat/internal/client/websocket_until_test.go

@@ -0,0 +1,103 @@
+package client
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/gorilla/websocket"
+	"github.com/stretchr/testify/require"
+)
+
+func TestClient_ReadFrameUntil_UsesCallerDeadlineInsteadOfRequestTimeout(t *testing.T) {
+	// Given
+	upgrader := websocket.Upgrader{}
+	serverReady := make(chan struct{})
+	allowFrame := make(chan struct{})
+	server := newWebSocketTestServer(t, upgrader, func(connection *websocket.Conn, _ *http.Request) {
+		close(serverReady)
+		<-allowFrame
+		require.NoError(t, connection.WriteMessage(websocket.BinaryMessage, []byte("held-session")))
+	})
+	client := newTestClient(t, Config{BaseURL: server.URL, RequestTimeout: 100 * time.Millisecond, MaxResponseBytes: 1024})
+	connection, err := client.DialWebSocket(context.Background(), "/held")
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = connection.Close() })
+	callerContext, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+	result := make(chan struct {
+		frame Frame
+		err   error
+	}, 1)
+
+	// When
+	go func() {
+		frame, readErr := connection.ReadFrameUntil(callerContext)
+		result <- struct {
+			frame Frame
+			err   error
+		}{frame: frame, err: readErr}
+	}()
+	<-serverReady
+	requestTimeout, cancelRequestTimeout := context.WithTimeout(context.Background(), 100*time.Millisecond)
+	defer cancelRequestTimeout()
+	<-requestTimeout.Done()
+	close(allowFrame)
+
+	// Then
+	select {
+	case readResult := <-result:
+		require.NoError(t, readResult.err)
+		require.Equal(t, FrameBinary, readResult.frame.Type)
+		require.Equal(t, []byte("held-session"), readResult.frame.Payload)
+	case <-time.After(time.Second):
+		t.Fatal("ReadFrameUntil did not receive the channel-released frame")
+	}
+}
+
+func TestClient_ReadFrameUntil_ReturnsParentCancellationAndUnblocksRead(t *testing.T) {
+	// Given
+	upgrader := websocket.Upgrader{}
+	serverReady := make(chan struct{})
+	server := newWebSocketTestServer(t, upgrader, func(_ *websocket.Conn, request *http.Request) {
+		close(serverReady)
+		<-request.Context().Done()
+	})
+	client := newTestClient(t, Config{BaseURL: server.URL, RequestTimeout: time.Second, MaxResponseBytes: 1024})
+	connection, err := client.DialWebSocket(context.Background(), "/cancel")
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = connection.Close() })
+	callerContext, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	result := make(chan error, 1)
+
+	// When
+	go func() {
+		_, readErr := connection.ReadFrameUntil(callerContext)
+		result <- readErr
+	}()
+	<-serverReady
+	cancel()
+
+	// Then
+	select {
+	case readErr := <-result:
+		require.ErrorIs(t, readErr, context.Canceled)
+	case <-time.After(time.Second):
+		t.Fatal("ReadFrameUntil did not stop after parent cancellation")
+	}
+}
+
+func newWebSocketTestServer(t *testing.T, upgrader websocket.Upgrader, serve func(*websocket.Conn, *http.Request)) *httptest.Server {
+	t.Helper()
+	server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
+		connection, err := upgrader.Upgrade(writer, request, nil)
+		require.NoError(t, err)
+		defer connection.Close()
+		serve(connection, request)
+	}))
+	t.Cleanup(server.Close)
+	return server
+}