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