woggioni
e4e00762bb
- kaya-core: WebSocket ABC gains session attribute and accept(headers=...) - kaya-core: AsgiWebSocket injects headers into websocket.accept message - kaya-rsgi: RsgiWebSocket accepts headers param (ignored — Granian's accept() takes no args) - kaya-session: SessionWebSocket wrapper exposes ws.session and injects Set-Cookie on accept() - kaya-session: SessionMixin registers before/after websocket hooks; session loaded at connect, persisted on close if modified - 10 new WV session tests covering read, persist, handshake cookie, regenerate, invalidate, isolation - Example and README updated
134 lines
4.6 KiB
Python
134 lines
4.6 KiB
Python
import unittest
|
|
import httpx
|
|
from pwo import async_test
|
|
from httpx_ws import aconnect_ws, WebSocketDisconnect
|
|
from httpx_ws.transport import ASGIWebSocketTransport
|
|
from kaya.core import KayaApp, WebSocket
|
|
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(WebSocketDisconnect) 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])
|