81 lines
2.9 KiB
Python
81 lines
2.9 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
|
|
|
|
|
|
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
|