nezha_domains/integration/agentcompat/internal/client/websocket_until_test.go

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
}