feat: Add async websocket_connect to AsyncTestClient (#3328)

feat: Add `websocket_connect` method to AsyncTestClient

Co-authored-by: kedod <kedod>

litestar/testing/client/async_client.py

@@ -1,14 +1,15 @@
 from __future__ import annotations
 
 from contextlib import AsyncExitStack
-from typing import TYPE_CHECKING, Any, Generic, Mapping, TypeVar
+from typing import TYPE_CHECKING, Any, Generic, Mapping, Sequence, TypeVar
+from urllib.parse import urljoin
 
 from httpx import USE_CLIENT_DEFAULT, AsyncClient, Response
 
 from litestar import HttpMethod
 from litestar.testing.client.base import BaseTestClient
 from litestar.testing.life_span_handler import LifeSpanHandler
-from litestar.testing.transport import TestClientTransport
+from litestar.testing.transport import ConnectionUpgradeExceptionError, TestClientTransport
 from litestar.types import AnyIOBackend, ASGIApp
 
 if TYPE_CHECKING:
@@ -27,6 +28,7 @@
     from typing_extensions import Self
 
     from litestar.middleware.session.base import BaseBackendConfig
+    from litestar.testing.websocket_test_session import WebSocketTestSession
 
 
 T = TypeVar("T", bound=ASGIApp)
@@ -468,6 +470,59 @@ async def delete(
             extensions=None if extensions is None else dict(extensions),
         )
 
+    async def websocket_connect(
+        self,
+        url: str,
+        subprotocols: Sequence[str] | None = None,
+        params: QueryParamTypes | None = None,
+        headers: HeaderTypes | None = None,
+        cookies: CookieTypes | None = None,
+        auth: AuthTypes | UseClientDefault = USE_CLIENT_DEFAULT,
+        follow_redirects: bool | UseClientDefault = USE_CLIENT_DEFAULT,
+        timeout: TimeoutTypes | UseClientDefault = USE_CLIENT_DEFAULT,
+        extensions: Mapping[str, Any] | None = None,
+    ) -> WebSocketTestSession:
+        """Sends a GET request to establish a websocket connection.
+
+        Args:
+            url: Request URL.
+            subprotocols: Websocket subprotocols.
+            params: Query parameters.
+            headers: Request headers.
+            cookies: Request cookies.
+            auth: Auth headers.
+            follow_redirects: Whether to follow redirects.
+            timeout: Request timeout.
+            extensions: Dictionary of ASGI extensions.
+
+        Returns:
+            A `WebSocketTestSession <litestar.testing.WebSocketTestSession>` instance.
+        """
+        url = urljoin("ws://testserver", url)
+        default_headers: dict[str, str] = {}
+        default_headers.setdefault("connection", "upgrade")
+        default_headers.setdefault("sec-websocket-key", "testserver==")
+        default_headers.setdefault("sec-websocket-version", "13")
+        if subprotocols is not None:
+            default_headers.setdefault("sec-websocket-protocol", ", ".join(subprotocols))
+        try:
+            await AsyncClient.request(
+                self,
+                "GET",
+                url,
+                headers={**dict(headers or {}), **default_headers},  # type: ignore[misc]
+                params=params,
+                cookies=cookies,
+                auth=auth,
+                follow_redirects=follow_redirects,
+                timeout=timeout,
+                extensions=None if extensions is None else dict(extensions),
+            )
+        except ConnectionUpgradeExceptionError as exc:
+            return exc.session
+
+        raise RuntimeError("Expected WebSocket upgrade")  # pragma: no cover
+
     async def get_session_data(self) -> dict[str, Any]:
         """Get session data.
 

tests/unit/test_testing/test_test_client.py

@@ -5,7 +5,7 @@
 
 from litestar import Controller, WebSocket, delete, head, patch, put, websocket
 from litestar.status_codes import HTTP_200_OK, HTTP_201_CREATED, HTTP_204_NO_CONTENT
-from litestar.testing import AsyncTestClient, WebSocketTestSession, create_test_client
+from litestar.testing import AsyncTestClient, WebSocketTestSession, create_async_test_client, create_test_client
 
 if TYPE_CHECKING:
     from litestar.middleware.session.base import BaseBackendConfig
@@ -261,3 +261,52 @@ async def handler(socket: WebSocket) -> None:
         Empty
     ), client.websocket_connect("/"):
         pass
+
+
+@pytest.mark.parametrize("block,timeout", [(False, None), (False, 0.001), (True, 0.001)])
+@pytest.mark.parametrize(
+    "receive_method",
+    [
+        WebSocketTestSession.receive,
+        WebSocketTestSession.receive_json,
+        WebSocketTestSession.receive_text,
+        WebSocketTestSession.receive_bytes,
+    ],
+)
+async def test_websocket_test_session_block_timeout_async(
+    receive_method: Callable[..., Any], block: bool, timeout: Optional[float], anyio_backend: "AnyIOBackend"
+) -> None:
+    @websocket()
+    async def handler(socket: WebSocket) -> None:
+        await socket.accept()
+
+    with pytest.raises(Empty):
+        async with create_async_test_client(handler, backend=anyio_backend) as client:
+            with await client.websocket_connect("/") as ws:
+                receive_method(ws, timeout=timeout, block=block)
+
+
+async def test_websocket_accept_timeout_async(anyio_backend: "AnyIOBackend") -> None:
+    @websocket()
+    async def handler(socket: WebSocket) -> None:
+        pass
+
+    async with create_async_test_client(handler, backend=anyio_backend, timeout=0.1) as client:
+        with pytest.raises(Empty):
+            with await client.websocket_connect("/"):
+                pass
+
+
+async def test_websocket_connect_async(anyio_backend: "AnyIOBackend") -> None:
+    @websocket()
+    async def handler(socket: WebSocket) -> None:
+        await socket.accept()
+        data = await socket.receive_json()
+        await socket.send_json(data)
+        await socket.close()
+
+    async with create_async_test_client(handler, backend=anyio_backend, timeout=0.1) as client:
+        with await client.websocket_connect("/", subprotocols="wamp") as ws:
+            ws.send_json({"data": "123"})
+            data = ws.receive_json()
+            assert data == {"data": "123"}