''' IB quote restart regressions using the real Tractor asyncio bridge. ''' import asyncio from contextlib import asynccontextmanager as acm from types import SimpleNamespace from async_generator import aclosing from ib_async import Contract, IB, RequestError, Ticker import pytest import tractor import trio from piker.brokers.ib import feed from piker.brokers.ib import _util from piker.brokers.ib.api import Client, open_client_proxy async def exercise_quote_restart() -> list[str]: ''' Cancel an active quote consumer, then reopen in the same actor. Only the external IB connection is replaced. The production quote producer, eventkit callback, bridge, broadcaster and context cache all run unchanged in an asyncio-infected Tractor child. The farm OK event in stream_quotes() cancels its local scope after a reset; reproduce that cancellation after an actual ticker delivery. A cached owner can turn this into a normal bridge exit, feeding TrioTaskExited back into the consumer and killing brokerd. Two generations must close their IB subscriptions and asyncio tasks. Return callback failures separately so actor survival cannot be mistaken for clean teardown. This proves local restart semantics, not connectivity to a real Gateway. ''' contract = Contract(symbol='QQQ') tasks: list[asyncio.Task] = [] tickers: list[Ticker] = [] cancelled: list[Contract] = [] callback_errors: list[dict] = [] class ReplayIB: ''' Deliver one real eventkit ticker per subscription. ''' def reqMktData(self, con: Contract, opts: str) -> Ticker: ''' Emit once the production producer installs its callback. ''' task = asyncio.current_task() assert task is not None tasks.append(task) ticker = Ticker(contract=con) tickers.append(ticker) asyncio.get_running_loop().call_soon( ticker.updateEvent.emit, ticker, ) return ticker def cancelMktData(self, con: Contract) -> None: ''' Record the production producer's subscription teardown. ''' cancelled.append(con) client = SimpleNamespace(ib=ReplayIB()) @acm async def load_clients(**kwargs): ''' Supply an offline IB peer without opening broker connections. ''' yield {'ib.replay': client} loop = asyncio.get_running_loop() previous_handler = loop.get_exception_handler() loop.set_exception_handler( lambda loop, context: callback_errors.append(context) ) try: with pytest.MonkeyPatch.context() as patch: patch.setattr(feed, 'load_aio_clients', load_clients) patch.setattr( feed, 'get_preferred_data_client', lambda clients: ('ib.replay', client), ) for generation in range(2): with trio.CancelScope() as scope: async with feed.open_aio_quote_stream( 'qqq.nasdaq', contract=contract, ) as quotes: async with aclosing(quotes): ticker = await quotes.receive() assert ticker is tickers[generation] scope.cancel() await trio.lowlevel.checkpoint() assert scope.cancelled_caught assert len(cancelled) == generation + 1 assert tasks[generation].done() assert len(tickers[generation].updateEvent) == 0 # Flush done callbacks before accepting a clean restart. await tractor.to_asyncio.run_task(asyncio.sleep, delay=0) return [ repr(context['exception']) for context in callback_errors ] finally: loop.set_exception_handler(previous_handler) @pytest.mark.parametrize( 'check_callbacks', [False, True], ids=['actor-survival', 'callback-cleanliness'], ) def test_quote_restart_preserves_actor( reg_addr, check_callbacks: bool, ) -> None: ''' A farm-style quote restart must leave its broker actor usable. Use a real TCP actor nursery and an asyncio-infected child. Run two subscription generations, then invoke another RPC on the same actor. Both cases reject a propagated bridge exception, leaked quote task or lost actor. The strict case also rejects asyncio callback failures; the survival case reports those independently. The backout can restore restart survival while the strict case still exposes Tractor's completion-callback defect. ''' async def main() -> None: with trio.fail_after(10): async with ( tractor.open_root_actor(registry_addrs=[reg_addr]), tractor.open_nursery() as nursery, ): portal = await nursery.start_actor( 'ib_quote_replay', enable_modules=[__name__], infect_asyncio=True, ) errors = await portal.run(exercise_quote_restart) errors += await portal.run(exercise_quote_restart) await portal.cancel_actor() if check_callbacks: assert not errors, errors trio.run(main)