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 ConnectionRejected
from wsproto.utilities import LocalProtocolError

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 ``LocalProtocolError`` 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.

    '''
    attempts: int = 0

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

        raise LocalProtocolError('invalid websocket response')
        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(LocalProtocolError):
            async with _web_bs.open_autorecon_ws(
                url='wss://reject.invalid/ws',
                connect_timeout=0.1,
            ):
                pytest.fail(
                    'failed websocket unexpectedly connected'
                )

    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 == []