test(agentcompat): add protocol clients

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

integration/agentcompat/internal/client/websocket_close_test.go

@@ -0,0 +1,169 @@
+package client
+
+import (
+	"context"
+	"errors"
+	"net"
+	"net/http"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/gorilla/websocket"
+	"github.com/stretchr/testify/require"
+)
+
+func TestClient_WebSocketClose_RetainsConcurrentPhysicalCloseResult(t *testing.T) {
+	// Given
+	closeErr := errors.New("physical close failed")
+	recordedConnection := newRetainedCloseConn(closeErr)
+	server := newWebSocketTestServer(t, websocket.Upgrader{}, func(_ *websocket.Conn, request *http.Request) {
+		<-request.Context().Done()
+	})
+	connection := dialRetainedCloseWebSocket(t, server.URL, recordedConnection)
+
+	// When
+	results := make(chan error, 8)
+	for range cap(results) {
+		go func() { results <- connection.Close() }()
+	}
+
+	// Then
+	for range cap(results) {
+		require.Equal(t, closeErr, <-results)
+	}
+	require.Equal(t, 1, recordedConnection.closeCount())
+	require.Equal(t, closeErr, connection.Close())
+}
+
+func TestClient_WebSocketClose_RetainsReadCancellationCloseResult(t *testing.T) {
+	// Given
+	closeErr := errors.New("read cancellation close failed")
+	recordedConnection := newRetainedCloseConn(closeErr)
+	serverReady := make(chan struct{})
+	server := newWebSocketTestServer(t, websocket.Upgrader{}, func(_ *websocket.Conn, request *http.Request) {
+		close(serverReady)
+		<-request.Context().Done()
+	})
+	connection := dialRetainedCloseWebSocket(t, server.URL, recordedConnection)
+	readContext, cancel := context.WithCancel(context.Background())
+	readResult := make(chan error, 1)
+	go func() {
+		_, readErr := connection.ReadFrameUntil(readContext)
+		readResult <- readErr
+	}()
+	<-serverReady
+
+	// When
+	cancel()
+
+	// Then
+	require.ErrorIs(t, <-readResult, context.Canceled)
+	require.Equal(t, closeErr, connection.Close())
+	require.Equal(t, 1, recordedConnection.closeCount())
+}
+
+func TestClient_WebSocketClose_RetainsWriteCancellationCloseResult(t *testing.T) {
+	// Given
+	closeErr := errors.New("write cancellation close failed")
+	recordedConnection := newRetainedCloseConn(closeErr)
+	serverReady := make(chan struct{})
+	server := newWebSocketTestServer(t, websocket.Upgrader{}, func(_ *websocket.Conn, request *http.Request) {
+		close(serverReady)
+		<-request.Context().Done()
+	})
+	connection := dialRetainedCloseWebSocket(t, server.URL, recordedConnection)
+	<-serverReady
+	recordedConnection.blockWrites()
+	writeContext, cancel := context.WithCancel(context.Background())
+	writeResult := make(chan error, 1)
+	go func() {
+		writeResult <- connection.WriteFrame(writeContext, Frame{Type: FrameBinary, Payload: []byte("blocked")})
+	}()
+	recordedConnection.awaitWrite(t)
+
+	// When
+	cancel()
+
+	// Then
+	require.ErrorIs(t, <-writeResult, context.Canceled)
+	require.Equal(t, closeErr, connection.Close())
+	require.Equal(t, 1, recordedConnection.closeCount())
+}
+
+type retainedCloseConn struct {
+	net.Conn
+	mu              sync.Mutex
+	closeErr        error
+	closeCountValue int
+	writeBlocked    bool
+	writeEntered    chan struct{}
+	writeReleased   chan struct{}
+	releaseWrite    sync.Once
+}
+
+func newRetainedCloseConn(closeErr error) *retainedCloseConn {
+	return &retainedCloseConn{closeErr: closeErr, writeEntered: make(chan struct{}, 1), writeReleased: make(chan struct{})}
+}
+
+func (connection *retainedCloseConn) Write(payload []byte) (int, error) {
+	connection.mu.Lock()
+	blocked := connection.writeBlocked
+	connection.mu.Unlock()
+	if blocked {
+		connection.writeEntered <- struct{}{}
+		<-connection.writeReleased
+	}
+	return connection.Conn.Write(payload)
+}
+
+func (connection *retainedCloseConn) Close() error {
+	connection.mu.Lock()
+	connection.closeCountValue++
+	connection.mu.Unlock()
+	underlyingErr := connection.Conn.Close()
+	connection.releaseWrite.Do(func() { close(connection.writeReleased) })
+	if connection.closeErr != nil {
+		return connection.closeErr
+	}
+	return underlyingErr
+}
+
+func (connection *retainedCloseConn) blockWrites() {
+	connection.mu.Lock()
+	defer connection.mu.Unlock()
+	connection.writeBlocked = true
+}
+
+func (connection *retainedCloseConn) awaitWrite(t *testing.T) {
+	t.Helper()
+	select {
+	case <-connection.writeEntered:
+	case <-time.After(time.Second):
+		t.Fatal("WebSocket client did not enter the blocked write")
+	}
+}
+
+func (connection *retainedCloseConn) closeCount() int {
+	connection.mu.Lock()
+	defer connection.mu.Unlock()
+	return connection.closeCountValue
+}
+
+func dialRetainedCloseWebSocket(t *testing.T, baseURL string, recordedConnection *retainedCloseConn) *WebSocketConnection {
+	t.Helper()
+	dialer := *websocket.DefaultDialer
+	dialer.NetDialContext = func(ctx context.Context, network, address string) (net.Conn, error) {
+		connection, err := (&net.Dialer{}).DialContext(ctx, network, address)
+		if err != nil {
+			return nil, err
+		}
+		recordedConnection.Conn = connection
+		return recordedConnection, nil
+	}
+	client := newTestClient(t, Config{BaseURL: baseURL, RequestTimeout: time.Second, MaxResponseBytes: 1024, WebSocketDialer: &dialer})
+	connection, err := client.DialWebSocket(context.Background(), "/retained-close")
+	require.NoError(t, err)
+	t.Cleanup(func() { require.ErrorIs(t, connection.Close(), recordedConnection.closeErr) })
+	return connection
+}