piker/tests/test_ib_quote_restart.py

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