305 lines
9.5 KiB
Python
305 lines
9.5 KiB
Python
'''
|
|
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
|