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.
151 lines
4.9 KiB
Python
151 lines
4.9 KiB
Python
import unittest
|
|
from types import SimpleNamespace
|
|
from typing import Any
|
|
|
|
from granian.rsgi import ProtocolClosed
|
|
from pwo import async_test
|
|
|
|
from kaya.core import WebSocketDisconnect
|
|
from kaya.rsgi import RsgiWebSocket
|
|
|
|
|
|
class FakeScope:
|
|
scheme = 'ws'
|
|
path = '/ws'
|
|
query_string = ''
|
|
headers: dict[str, list[str]] = {}
|
|
client = '127.0.0.1:12345'
|
|
server = '127.0.0.1:80'
|
|
|
|
|
|
class FakeTransport:
|
|
"""A transport whose peer vanished without a close handshake."""
|
|
|
|
def __init__(self) -> None:
|
|
self.receive_calls = 0
|
|
self.sent: list[Any] = []
|
|
|
|
async def receive(self) -> Any:
|
|
self.receive_calls += 1
|
|
raise ProtocolClosed()
|
|
|
|
async def send_str(self, data: str) -> None:
|
|
self.sent.append(data)
|
|
raise ProtocolClosed()
|
|
|
|
async def send_bytes(self, data: bytes) -> None:
|
|
self.sent.append(data)
|
|
raise ProtocolClosed()
|
|
|
|
|
|
class CleanCloseTransport(FakeTransport):
|
|
"""A transport that reports a close message (kind 0)."""
|
|
|
|
async def receive(self) -> Any:
|
|
self.receive_calls += 1
|
|
return SimpleNamespace(kind=0, data=None)
|
|
|
|
|
|
class FakeProtocol:
|
|
def __init__(self, transport: FakeTransport) -> None:
|
|
self._transport = transport
|
|
self.closed: list[int] = []
|
|
|
|
async def accept(self) -> FakeTransport:
|
|
return self._transport
|
|
|
|
def close(self, code: int) -> None:
|
|
self.closed.append(code)
|
|
|
|
|
|
class ClosedOnAcceptProtocol(FakeProtocol):
|
|
"""A protocol whose peer vanished during the handshake."""
|
|
|
|
async def accept(self) -> FakeTransport:
|
|
raise ProtocolClosed()
|
|
|
|
|
|
class RsgiWebSocketTest(unittest.TestCase):
|
|
|
|
def test_misconfigured_granian(self):
|
|
class FakeProtocolWithoutWebsockets:
|
|
pass
|
|
|
|
with self.assertRaises(RuntimeError) as ctx:
|
|
RsgiWebSocket(FakeScope(), FakeProtocolWithoutWebsockets()) # type: ignore[arg-type]
|
|
|
|
self.assertIn('Granian was not configured for websockets', str(ctx.exception))
|
|
|
|
async def _socket(self, transport: FakeTransport) -> tuple[RsgiWebSocket, FakeProtocol]:
|
|
protocol = FakeProtocol(transport)
|
|
ws = RsgiWebSocket(FakeScope(), protocol) # type: ignore[arg-type]
|
|
await ws.accept()
|
|
return ws, protocol
|
|
|
|
@async_test
|
|
async def test_receive_returns_close_when_transport_is_gone(self):
|
|
ws, _ = await self._socket(FakeTransport())
|
|
message = await ws.receive()
|
|
self.assertEqual('close', message.kind)
|
|
|
|
@async_test
|
|
async def test_iteration_ends_when_transport_is_gone(self):
|
|
ws, _ = await self._socket(FakeTransport())
|
|
messages = [message async for message in ws]
|
|
self.assertEqual([], messages)
|
|
|
|
@async_test
|
|
async def test_receive_remembers_the_close(self):
|
|
transport = FakeTransport()
|
|
ws, _ = await self._socket(transport)
|
|
self.assertEqual('close', (await ws.receive()).kind)
|
|
self.assertEqual('close', (await ws.receive()).kind)
|
|
self.assertEqual(1, transport.receive_calls)
|
|
|
|
@async_test
|
|
async def test_send_text_raises_disconnect_when_transport_is_gone(self):
|
|
ws, _ = await self._socket(FakeTransport())
|
|
with self.assertRaises(WebSocketDisconnect) as ctx:
|
|
await ws.send_text('hello')
|
|
self.assertEqual(1006, ctx.exception.code)
|
|
|
|
@async_test
|
|
async def test_send_bytes_raises_disconnect_when_transport_is_gone(self):
|
|
ws, _ = await self._socket(FakeTransport())
|
|
with self.assertRaises(WebSocketDisconnect) as ctx:
|
|
await ws.send_bytes(b'hello')
|
|
self.assertEqual(1006, ctx.exception.code)
|
|
|
|
@async_test
|
|
async def test_sends_after_disconnect_do_not_touch_the_transport(self):
|
|
transport = FakeTransport()
|
|
ws, _ = await self._socket(transport)
|
|
with self.assertRaises(WebSocketDisconnect):
|
|
await ws.send_text('first')
|
|
with self.assertRaises(WebSocketDisconnect):
|
|
await ws.send_text('second')
|
|
self.assertEqual(['first'], transport.sent)
|
|
|
|
@async_test
|
|
async def test_send_after_a_clean_close_raises_disconnect(self):
|
|
transport = CleanCloseTransport()
|
|
ws, _ = await self._socket(transport)
|
|
self.assertEqual('close', (await ws.receive()).kind)
|
|
with self.assertRaises(WebSocketDisconnect):
|
|
await ws.send_text('hello')
|
|
self.assertEqual([], transport.sent)
|
|
|
|
@async_test
|
|
async def test_accept_raises_disconnect_when_transport_is_gone(self):
|
|
ws = RsgiWebSocket(FakeScope(), ClosedOnAcceptProtocol(FakeTransport())) # type: ignore[arg-type]
|
|
with self.assertRaises(WebSocketDisconnect) as ctx:
|
|
await ws.accept()
|
|
self.assertEqual(1006, ctx.exception.code)
|
|
|
|
@async_test
|
|
async def test_close_is_noop_once_the_transport_is_gone(self):
|
|
ws, protocol = await self._socket(FakeTransport())
|
|
await ws.receive()
|
|
await ws.close(1000)
|
|
self.assertEqual([], protocol.closed)
|