kaya/packages/kaya-core/tests/test_websocket.py

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])

    @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)