''' 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, 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. ''' 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)