proto-socket/python/test/test_ws.py

from __future__ import annotations

import asyncio
import ssl
from pathlib import Path

import pytest

from proto_socket.communicator import type_name_of
from proto_socket.packets.message_common_pb2 import TestData as ProtoTestData
from proto_socket.ws_client import WsClient, connect_ws, connect_wss
from proto_socket.ws_server import WsServer


def _make_server_ssl_context() -> ssl.SSLContext:
    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
    certs_dir = Path(__file__).parent / "certs"
    ctx.load_cert_chain(certs_dir / "server.crt", certs_dir / "server.key")
    return ctx


def _make_client_ssl_context() -> ssl.SSLContext:
    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
    certs_dir = Path(__file__).parent / "certs"
    ctx.load_verify_locations(certs_dir / "server.crt")
    ctx.check_hostname = False
    return ctx


def parser_map():
    return {type_name_of(ProtoTestData): ProtoTestData.FromString}


async def _wait_for_clients(server: WsServer, count: int = 1, timeout: float = 1.0) -> None:
    loop = asyncio.get_running_loop()
    deadline = loop.time() + timeout
    while loop.time() < deadline:
        if len(server.clients()) >= count:
            return
        await asyncio.sleep(0.01)
    raise TimeoutError("server did not see expected client count")


@pytest.mark.asyncio
async def test_ws_send_receive():
    received = asyncio.get_running_loop().create_future()
    server = WsServer(
        "127.0.0.1",
        0,
        "/",
        lambda ws: WsClient(ws, 0, 0, parser_map()),
    )
    server.on_client_connected = lambda client: client.communicator.add_listener(
        type_name_of(ProtoTestData),
        lambda message: received.set_result(message) if not received.done() else None,
    )
    await server.start()
    client = await connect_ws("127.0.0.1", server.port, "/", 0, 0, parser_map())

    await client.communicator.send(ProtoTestData(index=1, message="hello ws"))
    result = await asyncio.wait_for(received, 2)

    assert result.index == 1
    assert result.message == "hello ws"

    await client.close()
    await server.stop()


@pytest.mark.asyncio
async def test_ws_server_push():
    pushed = asyncio.get_running_loop().create_future()
    server = WsServer(
        "127.0.0.1",
        0,
        "/",
        lambda ws: WsClient(ws, 0, 0, parser_map()),
    )
    await server.start()
    client = await connect_ws("127.0.0.1", server.port, "/", 0, 0, parser_map())
    client.communicator.add_listener(
        type_name_of(ProtoTestData),
        lambda msg: pushed.set_result(msg) if not pushed.done() else None,
    )

    await _wait_for_clients(server)
    await server.broadcast(ProtoTestData(index=200, message="push from python ws server"))
    result = await asyncio.wait_for(pushed, 2)

    assert result.index == 200
    assert result.message == "push from python ws server"

    await client.close()
    await server.stop()


@pytest.mark.asyncio
async def test_ws_tls_send_receive():
    received = asyncio.get_running_loop().create_future()
    server = WsServer(
        "127.0.0.1",
        0,
        "/",
        lambda ws: WsClient(ws, 0, 0, parser_map()),
        ssl_context=_make_server_ssl_context(),
    )
    server.on_client_connected = lambda client: client.communicator.add_listener(
        type_name_of(ProtoTestData),
        lambda message: received.set_result(message) if not received.done() else None,
    )
    await server.start()
    client = await connect_wss(
        "127.0.0.1",
        server.port,
        "/",
        _make_client_ssl_context(),
        0,
        0,
        parser_map(),
    )

    await client.communicator.send(ProtoTestData(index=4, message="hello wss"))
    result = await asyncio.wait_for(received, 2)

    assert result.index == 4
    assert result.message == "hello wss"

    await client.close()
    await server.stop()


@pytest.mark.asyncio
async def test_ws_tls_request_response():
    server = WsServer(
        "127.0.0.1",
        0,
        "/",
        lambda ws: WsClient(ws, 0, 0, parser_map()),
        ssl_context=_make_server_ssl_context(),
    )
    server.on_client_connected = lambda client: client.communicator.add_request_listener(
        type_name_of(ProtoTestData),
        lambda req: ProtoTestData(index=req.index * 2, message="wss echo: " + req.message),
    )
    await server.start()
    client = await connect_wss(
        "127.0.0.1",
        server.port,
        "/",
        _make_client_ssl_context(),
        0,
        0,
        parser_map(),
    )

    result = await client.communicator.send_request(
        ProtoTestData(index=25, message="wss request"),
        ProtoTestData,
        timeout=2,
    )

    assert result.index == 50
    assert result.message == "wss echo: wss request"

    await client.close()
    await server.stop()


@pytest.mark.asyncio
async def test_ws_request_response():
    server = WsServer(
        "127.0.0.1",
        0,
        "/",
        lambda ws: WsClient(ws, 0, 0, parser_map()),
    )
    server.on_client_connected = lambda client: client.communicator.add_request_listener(
        type_name_of(ProtoTestData),
        lambda req: ProtoTestData(index=req.index * 2, message="echo: " + req.message),
    )
    await server.start()
    client = await connect_ws("127.0.0.1", server.port, "/", 0, 0, parser_map())

    result = await client.communicator.send_request(
        ProtoTestData(index=21, message="ws request"),
        ProtoTestData,
        timeout=2,
    )

    assert result.index == 42
    assert result.message == "echo: ws request"

    await client.close()
    await server.stop()


@pytest.mark.asyncio
async def test_ws_concurrent_requests():
    server = WsServer(
        "127.0.0.1",
        0,
        "/",
        lambda ws: WsClient(ws, 0, 0, parser_map()),
    )
    server.on_client_connected = lambda client: client.communicator.add_request_listener(
        type_name_of(ProtoTestData),
        lambda req: ProtoTestData(index=req.index * 2, message="echo: " + req.message),
    )
    await server.start()
    client = await connect_ws("127.0.0.1", server.port, "/", 0, 0, parser_map())

    async def do_request(i: int) -> None:
        index = 30 + i
        message = f"request {i}"
        result = await client.communicator.send_request(
            ProtoTestData(index=index, message=message),
            ProtoTestData,
            timeout=2,
        )
        assert result.index == index * 2
        assert result.message == f"echo: {message}"

    await asyncio.gather(*[do_request(i) for i in range(5)])
    await client.close()
    await server.stop()


@pytest.mark.asyncio
async def test_ws_inbound_backpressure_pauses_read():
    server_client = None
    
    def on_server_connect(client):
        nonlocal server_client
        server_client = client

    server = WsServer(
        "127.0.0.1",
        0,
        "/",
        lambda ws: WsClient(ws, 0, 0, parser_map()),
    )
    server.on_client_connected = on_server_connect
    await server.start()
    
    client = await connect_ws("127.0.0.1", server.port, "/", 0, 0, parser_map())
    await _wait_for_clients(server, count=1, timeout=2.0)
    
    handler_promise = asyncio.Event()

    async def blocking_handler(msg, nonce=0):
        await handler_promise.wait()
        return ProtoTestData(index=100)

    server_client.communicator.add_request_listener(type_name_of(ProtoTestData), blocking_handler)

    first_req_task = asyncio.create_task(
        client.communicator.send_request(
            ProtoTestData(index=1, message="first"),
            ProtoTestData,
            timeout=5.0,
        )
    )
    await asyncio.sleep(0.05)

    for i in range(2, 66):
        await client.communicator.send(ProtoTestData(index=i, message="heavy"))
        
    await client.communicator.send(ProtoTestData(index=66, message="overflow"))
    await asyncio.sleep(0.1)

    assert not server_client._read_task.done(), "ws read loop should not be finished"

    handler_promise.set()
    await asyncio.sleep(0.05)

    await first_req_task
    
    await client.close()
    await server.stop()