piker/tests/test_ems_broadcast.py

'''
EMS broadcast ownership regressions.

'''
from contextlib import asynccontextmanager as acm

import trio
from trio.testing import wait_all_tasks_blocked
from tractor.trionics import broadcast_receiver

from piker.clearing import _ems


def test_dark_clearing_owns_non_strict_subscription(
    monkeypatch,
) -> None:
    '''
    Dark clearing must not mutate another feed consumer's lag policy.

    `clear_dark_triggers()` previously assigned `_raise_on_lag` on
    whichever root or temporary receiver was stored in `Flume.stream`.
    That changed private shared state and made behavior depend on feed
    cache ownership. Supply a stream whose subscription records the
    public lag-policy argument and yields a distinct child. Replace
    the hot trigger loop with a checkpointing probe, then prove the
    non-strict child is entered, passed to the core, and closed only
    after processing completes while the source remains untouched.

    '''
    events: list[tuple] = []
    child = object()

    class QuoteStream:
        '''
        Record public subscription ownership without quote traffic.

        '''
        @acm
        async def subscribe(self, raise_on_lag: bool = True):
            '''
            Yield one dedicated child for the wrapper lifetime.

            '''
            events.append(('enter', raise_on_lag))
            try:
                yield child
            finally:
                events.append(('exit', raise_on_lag))

    async def clear_core(*args) -> None:
        '''
        Prove the child remains owned while processing runs.

        '''
        events.append(('core', args[2] is child))
        await trio.lowlevel.checkpoint()

    monkeypatch.setattr(
        _ems,
        '_clear_dark_triggers',
        clear_core,
    )

    async def main() -> None:
        stream = QuoteStream()
        await _ems.clear_dark_triggers(
            object(),
            object(),
            stream,
            'ib',
            'nvda.nasdaq.ib',
            object(),
        )
        assert not hasattr(stream, '_raise_on_lag')

    trio.run(main)
    assert events == [
        ('enter', False),
        ('core', True),
        ('exit', False),
    ]


def test_dark_clearing_with_real_broadcast() -> None:
    '''
    Dark clearing must consume quotes without owning sibling lifetime.

    Exercise the production trigger loop and real Tractor broadcaster
    together. Deliver one trade tick to both EMS and a sibling cursor;
    DarkBook.lasts proves EMS processed it. Cancel only the EMS task
    and verify its subscription disappears while the sibling still
    receives the next tick with its strict lag policy intact.

    The book has no order triggers, so this requires no broker or
    account. Unlike the imported wrapper test, no stream or clearing
    function is replaced.

    '''
    async def main() -> None:
        symbol = 'qqq.nasdaq.ib'
        book = _ems.DarkBook(broker='ib', triggers={}, lasts={})
        tx, rx = trio.open_memory_channel(8)
        root = broadcast_receiver(rx, 8)
        quote = {symbol: {'ticks': [{'type': 'trade', 'price': 100}]}}

        with trio.fail_after(2):
            async with tx, rx, root, root.subscribe() as sibling:
                async with trio.open_nursery() as nursery:
                    nursery.start_soon(
                        _ems.clear_dark_triggers,
                        None,
                        None,
                        root,
                        'ib',
                        symbol,
                        book,
                    )
                    await wait_all_tasks_blocked()
                    assert len(root._state.subs) == 3
                    await tx.send(quote)
                    assert await sibling.receive() == quote
                    await wait_all_tasks_blocked()
                    assert book.lasts[symbol] == 100
                    assert root._raise_on_lag
                    assert sibling._raise_on_lag
                    nursery.cancel_scope.cancel()

                assert len(root._state.subs) == 2
                assert not sibling._closed
                await tx.send(quote)
                assert await sibling.receive() == quote

    trio.run(main)