piker/tests/test_websocket.py

'''
Generic websocket connection and teardown regressions.

'''
from contextlib import asynccontextmanager as acm

import pytest
import trio
from trio_websocket import (
    ConnectionClosed,
    ConnectionRejected,
    WebSocketConnection,
    serve_websocket,
)
from wsproto.events import CloseConnection

from piker.data import _web_bs


def test_websocket_connection_startup_fails(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    '''
    ``open_autorecon_ws()`` previously had no bound while the real
    ``_reconnect_forever()`` task awaited its first websocket
    handshake, so provider startup inherited a 616-second deadline.

    Replace only ``_web_bs.open_websocket_url()`` with a context
    entry that never completes. The URL-bearing ``TimeoutError``
    proves the wrapper bounds the real reconnect task and cancels its
    nursery instead of waiting indefinitely for
    ``NoBsWs._connected``.

    '''
    @acm
    async def never_connect(url: str):
        await trio.sleep_forever()
        yield

    monkeypatch.setattr(
        _web_bs,
        'open_websocket_url',
        never_connect,
    )

    async def main() -> None:
        with pytest.raises(
            TimeoutError,
            match='wss://silent.invalid/ws',
        ):
            async with _web_bs.open_autorecon_ws(
                url='wss://silent.invalid/ws',
                connect_timeout=0.01,
            ):
                pytest.fail(
                    'silent websocket unexpectedly connected'
                )

    trio.run(main)


def test_immediate_websocket_failure_retries(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    '''
    ``_reconnect_forever()`` previously fell through after a handled
    first-handshake rejection and read ``cs`` before any websocket
    nursery had assigned it, masking the transport failure with an
    unbound-local error.

    Make the real reconnect loop receive ``ConnectionRejected`` on
    its first context entry and ``RuntimeError`` on its second.
    Removing retry sleeps fixes their order. Two attempts and the
    propagated second error prove the first rejection restarts the
    loop without touching connection-local cancellation state.
    Require the exact error object so wrapping cannot discard its
    existing diagnostic context.

    '''
    attempts: int = 0
    failure = RuntimeError('unexpected websocket failure')

    @acm
    async def reject_url(url: str):
        nonlocal attempts
        attempts += 1
        if attempts == 1:
            raise ConnectionRejected(503, [], b'unavailable')

        raise failure
        yield

    async def no_throttle(delay: float) -> None:
        await trio.lowlevel.checkpoint()

    monkeypatch.setattr(
        _web_bs,
        'open_websocket_url',
        reject_url,
    )
    monkeypatch.setattr(
        _web_bs.trio,
        'sleep',
        no_throttle,
    )

    async def main() -> None:
        with pytest.raises(
            RuntimeError,
            match='unexpected websocket',
        ) as caught:
            async with _web_bs.open_autorecon_ws(
                url='wss://reject.invalid/ws',
                connect_timeout=0.1,
            ):
                pytest.fail(
                    'failed websocket unexpectedly connected'
                )
        assert caught.value is failure

    trio.run(main)
    assert attempts == 2


def test_reconnect_cancellation_bypasses_error_throttle(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    '''
    ``_reconnect_forever()`` previously caught ``trio`` cancellation
    in its broad ``BaseException`` handler and entered a shielded
    200ms error throttle, delaying websocket and feed-task teardown.

    Block the real reconnect task in the patched websocket context,
    wait until that cancellation point is active, and cancel its
    nursery. The tracked ``trio.sleep()`` replacement records any
    retry throttle without wall time. No calls after nursery exit
    proves structured cancellation bypasses normal error throttling.

    '''
    entered = trio.Event()
    sleep_calls: list[float] = []

    @acm
    async def blocked_url(url: str):
        entered.set()
        await trio.sleep_forever()
        yield

    async def record_sleep(delay: float) -> None:
        sleep_calls.append(delay)
        await trio.lowlevel.checkpoint()

    monkeypatch.setattr(
        _web_bs,
        'open_websocket_url',
        blocked_url,
    )
    monkeypatch.setattr(
        _web_bs.trio,
        'sleep',
        record_sleep,
    )

    async def main() -> None:
        send, recv = trio.open_memory_channel(1)
        nobsws = _web_bs.NoBsWs(
            'wss://cancel.invalid/ws',
            recv,
            16,
        )
        async with (
            send,
            recv,
            trio.open_nursery() as tn,
        ):
            await tn.start(
                _web_bs._reconnect_forever,
                'wss://cancel.invalid/ws',
                send,
                nobsws,
                3,
            )
            await entered.wait()
            tn.cancel_scope.cancel()

    trio.run(main)
    assert sleep_calls == []


@pytest.mark.parametrize('shutdown', ['recover', 'reader', 'retry'])
def test_reader_ping_close_recovery(
    monkeypatch: pytest.MonkeyPatch,
    caplog: pytest.LogCaptureFixture,
    shutdown: str,
) -> None:
    '''
    A real websocket reader's Pong error escaped reconnect ownership.

    Binance pings arriving during local closure raised wsproto's
    LocalProtocolError in trio-websocket's reader nursery, outside
    proxy_msgs(). 97eb8347 propagated that error from the outer
    wrapper, terminating the persistent broker instead of reconnecting.

    Use real loopback TCP, handshake, Ping frames and reader dispatch.
    At the first client Ping, synchronously advance its real wsproto
    connection to LOCAL_CLOSING. Withhold that Close frame from the
    wire so a peer Close reply cannot race this deterministic gate.
    The original Ping handler then generates the actual Pong error.
    Require a replacement connection to deliver fresh application data.

    For shutdown, cancel either at the held Ping boundary or after
    the actual error reaches the reconnect throttle. Require bounded
    exit and exactly one connection, proving teardown never starts a
    retry. Only the Ping handler ordering is controlled; contexts,
    readers, cancellation, reconnect tasks and delivery remain real.

    '''
    caplog.set_level('INFO', logger=_web_bs.log.name)

    async def main() -> None:
        connections: int = 0
        ping_entered = trio.Event()
        handler_exited = trio.Event()
        retry_entered = trio.Event()
        ping_handler = WebSocketConnection._handle_ping_event
        faulted: list[WebSocketConnection] = []
        log_exception = _web_bs.log.exception

        def observe_retry(message, *args, **kwargs) -> None:
            if message.startswith('Retrying connection:'):
                retry_entered.set()
            log_exception(message, *args, **kwargs)

        monkeypatch.setattr(_web_bs.log, 'exception', observe_retry)

        async def gate_ping(ws, event) -> None:
            if ws.is_client and not faulted:
                faulted.append(ws)
                ping_entered.set()
                try:
                    if shutdown == 'reader':
                        await trio.sleep_forever()
                    ws._wsproto.send(CloseConnection(code=1000))
                    await ping_handler(ws, event)
                finally:
                    handler_exited.set()
            else:
                await ping_handler(ws, event)

        monkeypatch.setattr(
            WebSocketConnection, '_handle_ping_event', gate_ping,
        )

        async def serve(request) -> None:
            nonlocal connections
            ws = await request.accept()
            connections += 1
            if connections == 1:
                try:
                    await ws.ping(b'close-race')
                except ConnectionClosed:
                    pass
            else:
                await ws.send_message('{"generation": 2}')
            await trio.sleep_forever()

        async with trio.open_nursery() as tn:
            server = await tn.start(
                serve_websocket, serve, '127.0.0.1', 0, None,
            )
            url = f'ws://127.0.0.1:{server.port}/'
            with trio.fail_after(5):
                with trio.CancelScope() as owner:
                    async with _web_bs.open_autorecon_ws(
                        url, connect_timeout=2,
                    ) as ws:
                        await ping_entered.wait()
                        if shutdown == 'retry':
                            await retry_entered.wait()
                        if shutdown != 'recover':
                            owner.cancel()
                            await trio.lowlevel.checkpoint()
                        else:
                            assert await ws.recv_msg() == {
                                'generation': 2,
                            }
                await handler_exited.wait()
                assert len(faulted) == 1
                assert connections == (
                    2 if shutdown == 'recover' else 1
                )
            tn.cancel_scope.cancel()

    trio.run(main)
    if shutdown == 'recover':
        messages = [record.getMessage() for record in caplog.records]
        assert sum(
            message.startswith('Websocket ready:')
            for message in messages
        ) == 2
        assert sum(
            message.startswith('Websocket first message relayed:')
            for message in messages
        ) == 1