woggioni
7c46badebb
CI / Build Pip package (push) Successful in 4m55s
Granian's RSGI transport raises RSGIProtocolClosed when the peer is gone without a close handshake, which escaped receive() and was logged as an application error, and sends to the dead transport raised the raw Granian exception as well. Treat a closed RSGI transport like the ASGI websocket.disconnect message: receive() reports a close message so async-for loops end normally, while writes (accept/send/close) raise a new kaya.core.WebSocketDisconnect carrying close code 1006. KayaApp.handle_websocket swallows that exception as normal termination, so after hooks still run and disconnects are no longer reported as errors. The ASGI adapter maps its server's OSError disconnect (e.g. uvicorn's ClientDisconnected) to the same exception for consistent semantics across backends.
264 lines
8.7 KiB
Python
264 lines
8.7 KiB
Python
import unittest
|
|
import httpx
|
|
from pwo import async_test
|
|
from httpx_ws import aconnect_ws, WebSocketDisconnect as HttpxWebSocketDisconnect
|
|
from httpx_ws.transport import ASGIWebSocketTransport
|
|
from kaya.core import KayaApp, WebSocket, WebSocketDisconnect, WebSocketMessage
|
|
from kaya.core._asgi import AsgiWebSocket
|
|
|
|
|
|
class WebSocketTest(unittest.TestCase):
|
|
app: KayaApp
|
|
|
|
def setUp(self):
|
|
self.app = KayaApp()
|
|
|
|
@self.app.websocket('/echo')
|
|
async def echo(ws: WebSocket) -> None:
|
|
await ws.accept()
|
|
async for msg in ws:
|
|
if msg.kind == 'text':
|
|
await ws.send_text(f"echo: {msg.data}")
|
|
elif msg.kind == 'binary':
|
|
data = msg.data
|
|
assert isinstance(data, bytes)
|
|
await ws.send_bytes(data)
|
|
|
|
@self.app.websocket('/room/${room_id}')
|
|
async def room(ws: WebSocket, room_id: str) -> None:
|
|
await ws.accept()
|
|
async for msg in ws:
|
|
if msg.kind == 'text':
|
|
await ws.send_text(f"[{room_id}] {msg.data}")
|
|
|
|
@async_test
|
|
async def test_echo_text(self):
|
|
transport = ASGIWebSocketTransport(app=self.app)
|
|
|
|
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):
|
|
transport = ASGIWebSocketTransport(app=self.app)
|
|
|
|
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):
|
|
transport = ASGIWebSocketTransport(app=self.app)
|
|
|
|
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):
|
|
transport = ASGIWebSocketTransport(app=self.app)
|
|
|
|
async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as client:
|
|
with self.assertRaises(HttpxWebSocketDisconnect) as cm:
|
|
async with aconnect_ws("/unknown", client) as ws:
|
|
pass
|
|
|
|
self.assertEqual(cm.exception.code, 1000)
|
|
|
|
@async_test
|
|
async def test_close_from_client(self):
|
|
transport = ASGIWebSocketTransport(app=self.app)
|
|
|
|
async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as client:
|
|
async with aconnect_ws("/echo", client) as ws:
|
|
pass
|
|
|
|
@async_test
|
|
async def test_accept_with_headers(self):
|
|
sent_messages = []
|
|
|
|
async def send(message):
|
|
sent_messages.append(message)
|
|
|
|
async def receive():
|
|
return {'type': 'websocket.connect'}
|
|
|
|
scope = {
|
|
'type': 'websocket',
|
|
'path': '/echo',
|
|
'query_string': b'',
|
|
'scheme': 'ws',
|
|
'client': ('127.0.0.1', 12345),
|
|
'server': ('127.0.0.1', 80),
|
|
'headers': [],
|
|
}
|
|
ws = AsgiWebSocket(scope, receive, send)
|
|
await ws.accept(headers={'Set-Cookie': 'sid=abc; Path=/', 'X-Custom': ('a', 'b')})
|
|
self.assertEqual(1, len(sent_messages))
|
|
message = sent_messages[0]
|
|
self.assertEqual('websocket.accept', message['type'])
|
|
self.assertIn((b'Set-Cookie', b'sid=abc; Path=/'), message['headers'])
|
|
self.assertIn((b'X-Custom', b'a'), message['headers'])
|
|
self.assertIn((b'X-Custom', b'b'), message['headers'])
|
|
|
|
@async_test
|
|
async def test_accept_without_headers(self):
|
|
sent_messages = []
|
|
|
|
async def send(message):
|
|
sent_messages.append(message)
|
|
|
|
async def receive():
|
|
return {'type': 'websocket.connect'}
|
|
|
|
scope = {
|
|
'type': 'websocket',
|
|
'path': '/echo',
|
|
'query_string': b'',
|
|
'scheme': 'ws',
|
|
'client': ('127.0.0.1', 12345),
|
|
'server': ('127.0.0.1', 80),
|
|
'headers': [],
|
|
}
|
|
ws = AsgiWebSocket(scope, receive, send)
|
|
await ws.accept()
|
|
self.assertEqual(1, len(sent_messages))
|
|
self.assertEqual({'type': 'websocket.accept'}, sent_messages[0])
|
|
|
|
@async_test
|
|
async def test_websocket_scope_without_scheme(self):
|
|
# Daphne omits the optional `scheme` key from websocket scopes.
|
|
async def send(message):
|
|
pass
|
|
|
|
async def receive():
|
|
return {'type': 'websocket.connect'}
|
|
|
|
scope = {
|
|
'type': 'websocket',
|
|
'path': '/echo',
|
|
'query_string': b'',
|
|
'client': ('127.0.0.1', 12345),
|
|
'server': ('127.0.0.1', 80),
|
|
'headers': [],
|
|
}
|
|
ws = AsgiWebSocket(scope, receive, send)
|
|
self.assertEqual('ws', ws.scheme)
|
|
|
|
|
|
class FakeWebSocket(WebSocket):
|
|
"""Minimal in-memory websocket for exercising ``KayaApp`` directly."""
|
|
|
|
def __init__(self, path: str) -> None:
|
|
self.path = path
|
|
self.query_string = ''
|
|
self.scheme = 'ws'
|
|
self.headers: dict[str, list[str]] = {}
|
|
self.client = ('127.0.0.1', 12345)
|
|
self.server = ('127.0.0.1', 80)
|
|
self.closed: list[int] = []
|
|
|
|
async def accept(self, headers=None) -> None:
|
|
pass
|
|
|
|
async def receive(self) -> WebSocketMessage:
|
|
return WebSocketMessage(kind='close')
|
|
|
|
async def send_text(self, data: str) -> None:
|
|
pass
|
|
|
|
async def send_bytes(self, data: bytes) -> None:
|
|
pass
|
|
|
|
async def close(self, code: int = 1000) -> None:
|
|
self.closed.append(code)
|
|
|
|
async def __anext__(self) -> WebSocketMessage:
|
|
raise StopAsyncIteration
|
|
|
|
|
|
class WebSocketDisconnectTest(unittest.TestCase):
|
|
|
|
def setUp(self):
|
|
self.app = KayaApp()
|
|
|
|
@self.app.websocket('/disconnect')
|
|
async def disconnect(ws: WebSocket) -> None:
|
|
await ws.accept()
|
|
await ws.send_text('before')
|
|
raise WebSocketDisconnect()
|
|
|
|
@self.app.websocket('/failure')
|
|
async def failure(ws: WebSocket) -> None:
|
|
await ws.accept()
|
|
raise RuntimeError('boom')
|
|
|
|
@async_test
|
|
async def test_disconnect_terminates_quietly(self):
|
|
after_hook_calls = []
|
|
|
|
async def after(ws: WebSocket) -> None:
|
|
after_hook_calls.append(ws)
|
|
|
|
self.app.add_after_websocket_hook(after)
|
|
ws = FakeWebSocket('/disconnect')
|
|
await self.app.handle_websocket(ws)
|
|
self.assertEqual([ws], after_hook_calls)
|
|
self.assertIsNone(ws.exception)
|
|
|
|
@async_test
|
|
async def test_other_errors_are_recorded_and_reraised(self):
|
|
ws = FakeWebSocket('/failure')
|
|
with self.assertRaises(RuntimeError):
|
|
await self.app.handle_websocket(ws)
|
|
self.assertIsInstance(ws.exception, RuntimeError)
|
|
|
|
@async_test
|
|
async def test_asgi_send_after_disconnect_raises_websocket_disconnect(self):
|
|
async def send(message):
|
|
raise OSError('client gone')
|
|
|
|
async def receive():
|
|
return {'type': 'websocket.disconnect', 'code': 1006}
|
|
|
|
ws = AsgiWebSocket(self._scope(), receive, send)
|
|
with self.assertRaises(WebSocketDisconnect) as ctx:
|
|
await ws.send_text('hello')
|
|
self.assertEqual(1006, ctx.exception.code)
|
|
with self.assertRaises(WebSocketDisconnect):
|
|
await ws.send_bytes(b'hello')
|
|
with self.assertRaises(WebSocketDisconnect):
|
|
await ws.close()
|
|
|
|
@async_test
|
|
async def test_asgi_accept_after_disconnect_raises_websocket_disconnect(self):
|
|
async def send(message):
|
|
raise OSError('client gone')
|
|
|
|
async def receive():
|
|
return {'type': 'websocket.connect'}
|
|
|
|
ws = AsgiWebSocket(self._scope(), receive, send)
|
|
with self.assertRaises(WebSocketDisconnect) as ctx:
|
|
await ws.accept()
|
|
self.assertEqual(1006, ctx.exception.code)
|
|
|
|
@staticmethod
|
|
def _scope() -> dict:
|
|
return {
|
|
'type': 'websocket',
|
|
'path': '/echo',
|
|
'query_string': b'',
|
|
'scheme': 'ws',
|
|
'client': ('127.0.0.1', 12345),
|
|
'server': ('127.0.0.1', 80),
|
|
'headers': [],
|
|
}
|