Refactor websocket tests to use httpx-ws ASGI transport

2026-07-15 14:33:03 +00:00
parent 0d6ecce415
commit 59f38f067f
3 changed files with 34 additions and 106 deletions
@@ -28,7 +28,7 @@ dependencies = [

 [project.optional-dependencies]
 dev = [
-    "build", "mypy", "ipdb", "twine", "httpx"
+    "build", "mypy", "ipdb", "twine", "httpx", "httpx-ws"
 ]

 [project.urls]
@@ -1,25 +1,9 @@
 import unittest
-from typing import Any, Callable, Awaitable, List, Mapping, Optional
+import httpx
 from pwo import async_test
-from kaya.core import KayaApp, WebSocket, WebSocketMessage
-
-
-def websocket_scope(path: str = '/ws') -> Mapping[str, Any]:
-    return {
-        'type': 'websocket',
-        'asgi': {'spec_version': '2.3', 'version': '3.0'},
-        'http_version': '1.1',
-        'scheme': 'ws',
-        'path': path,
-        'raw_path': path.encode(),
-        'query_string': b'',
-        'root_path': '',
-        'headers': [],
-        'client': ('127.0.0.1', 12345),
-        'server': ('127.0.0.1', 80),
-        'subprotocols': [],
-        'extensions': None,
-    }
+from httpx_ws import aconnect_ws, WebSocketDisconnect
+from httpx_ws.transport import ASGIWebSocketTransport
+from kaya.core import KayaApp, WebSocket


 class WebSocketTest(unittest.TestCase):
@@ -48,106 +32,49 @@ class WebSocketTest(unittest.TestCase):

    @async_test
    async def test_echo_text(self):
-        sent_messages: List[Mapping[str, Any]] = []
-        received_messages: List[Mapping[str, Any]] = []
+        transport = ASGIWebSocketTransport(app=self.app)

-        async def receive() -> Mapping[str, Any]:
-            if not received_messages:
-                received_messages.append({'type': 'websocket.connect'})
-                return received_messages[-1]
-            if len(received_messages) == 1:
-                received_messages.append({'type': 'websocket.receive', 'text': 'hello'})
-                return received_messages[-1]
-            received_messages.append({'type': 'websocket.disconnect', 'code': 1000})
-            return received_messages[-1]
-
-        async def send(message: Mapping[str, Any]) -> None:
-            sent_messages.append(message)
-
-        await self.app(websocket_scope('/echo'), receive, send)
-
-        self.assertEqual(sent_messages[0]['type'], 'websocket.accept')
-        self.assertEqual(sent_messages[1]['type'], 'websocket.send')
-        self.assertEqual(sent_messages[1]['text'], 'echo: hello')
+        async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as client:
+            async with aconnect_ws("/echo", client) as ws:
+                await ws.send_text("hello")
+                message = await ws.receive_text()
+                self.assertEqual(message, "echo: hello")

    @async_test
    async def test_echo_binary(self):
-        sent_messages: List[Mapping[str, Any]] = []
-        received_messages: List[Mapping[str, Any]] = []
+        transport = ASGIWebSocketTransport(app=self.app)

-        async def receive() -> Mapping[str, Any]:
-            if not received_messages:
-                received_messages.append({'type': 'websocket.connect'})
-                return received_messages[-1]
-            if len(received_messages) == 1:
-                received_messages.append({'type': 'websocket.receive', 'bytes': b'hello'})
-                return received_messages[-1]
-            received_messages.append({'type': 'websocket.disconnect', 'code': 1000})
-            return received_messages[-1]
-
-        async def send(message: Mapping[str, Any]) -> None:
-            sent_messages.append(message)
-
-        await self.app(websocket_scope('/echo'), receive, send)
-
-        self.assertEqual(sent_messages[0]['type'], 'websocket.accept')
-        self.assertEqual(sent_messages[1]['type'], 'websocket.send')
-        self.assertEqual(sent_messages[1]['bytes'], b'hello')
+        async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as client:
+            async with aconnect_ws("/echo", client) as ws:
+                await ws.send_bytes(b"hello")
+                message = await ws.receive_bytes()
+                self.assertEqual(message, b"hello")

    @async_test
    async def test_path_parameter(self):
-        sent_messages: List[Mapping[str, Any]] = []
-        received_messages: List[Mapping[str, Any]] = []
+        transport = ASGIWebSocketTransport(app=self.app)

-        async def receive() -> Mapping[str, Any]:
-            if not received_messages:
-                received_messages.append({'type': 'websocket.connect'})
-                return received_messages[-1]
-            if len(received_messages) == 1:
-                received_messages.append({'type': 'websocket.receive', 'text': 'hi'})
-                return received_messages[-1]
-            received_messages.append({'type': 'websocket.disconnect', 'code': 1000})
-            return received_messages[-1]
-
-        async def send(message: Mapping[str, Any]) -> None:
-            sent_messages.append(message)
-
-        await self.app(websocket_scope('/room/general'), receive, send)
-
-        self.assertEqual(sent_messages[0]['type'], 'websocket.accept')
-        self.assertEqual(sent_messages[1]['text'], '[general] hi')
+        async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as client:
+            async with aconnect_ws("/room/general", client) as ws:
+                await ws.send_text("hi")
+                message = await ws.receive_text()
+                self.assertEqual(message, "[general] hi")

    @async_test
    async def test_no_handler(self):
-        sent_messages: List[Mapping[str, Any]] = []
+        transport = ASGIWebSocketTransport(app=self.app)

-        async def receive() -> Mapping[str, Any]:
-            return {'type': 'websocket.connect'}
+        async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as client:
+            with self.assertRaises(WebSocketDisconnect) as cm:
+                async with aconnect_ws("/unknown", client) as ws:
+                    pass

-        async def send(message: Mapping[str, Any]) -> None:
-            sent_messages.append(message)
-
-        await self.app(websocket_scope('/unknown'), receive, send)
-
-        self.assertEqual(sent_messages[0]['type'], 'websocket.close')
-        self.assertEqual(sent_messages[0]['code'], 1000)
+            self.assertEqual(cm.exception.code, 1000)

    @async_test
    async def test_close_from_client(self):
-        sent_messages: List[Mapping[str, Any]] = []
-        received_messages: List[Mapping[str, Any]] = []
+        transport = ASGIWebSocketTransport(app=self.app)

-        async def receive() -> Mapping[str, Any]:
-            if not received_messages:
-                received_messages.append({'type': 'websocket.connect'})
-                return received_messages[-1]
-            received_messages.append({'type': 'websocket.disconnect', 'code': 1001})
-            return received_messages[-1]
-
-        async def send(message: Mapping[str, Any]) -> None:
-            sent_messages.append(message)
-
-        await self.app(websocket_scope('/echo'), receive, send)
-
-        self.assertEqual(sent_messages[0]['type'], 'websocket.accept')
-        self.assertEqual(len(sent_messages), 1)
+        async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as client:
+            async with aconnect_ws("/echo", client) as ws:
+                pass
@@ -11,6 +11,7 @@ mypy
 ipdb
 twine
 httpx
+httpx-ws
 granian
 pwo
 typing-extensions