piker/tests/test_feed_broadcast.py

'''
Cached feed broadcast ownership regressions.

'''
from contextlib import asynccontextmanager as acm
from types import SimpleNamespace

import tractor
import trio
import pytest

from piker.data import feed as feed_mod
from piker.data.flows import Flume


class FakeMsgStream:
    '''
    Duplex test stream backed by a real Tractor broadcaster.

    '''
    def __init__(self, broker: str) -> None:
        self.broker: str = broker
        self._ctx: object = object()
        self._tx, rx = trio.open_memory_channel(8)
        self._broadcaster = tractor.trionics.broadcast_receiver(
            rx,
            8,
        )
        self.children: list[
            tractor.trionics.BroadcastReceiver,
        ] = []
        self.controls: list[str] = []

    @property
    def ctx(self) -> object:
        '''
        Expose the public context ref provided by `MsgStream.ctx`.

        '''
        return self._ctx

    @acm
    async def subscribe(self):
        '''
        Yield a real child receiver with duplex send delegation.

        '''
        child: tractor.trionics.BroadcastReceiver
        async with self._broadcaster.subscribe() as child:
            child.send = self.send
            self.children.append(child)
            yield child

    async def send(self, msg: str) -> None:
        '''
        Record a pause or resume control sent through a child.

        '''
        self.controls.append(msg)

    async def push(self, msg: dict) -> None:
        '''
        Send one provider quote into the broadcast source.

        '''
        await self._tx.send(msg)

    async def close_source(self) -> None:
        '''
        End this provider without closing sibling relay channels.

        '''
        await self._tx.aclose()


def make_feed() -> tuple[
    feed_mod.Feed,
    dict[str, FakeMsgStream],
]:
    '''
    Build a real `Feed` and `Flume` graph for two providers.

    '''
    roots: dict[str, FakeMsgStream] = {
        broker: FakeMsgStream(broker)
        for broker in ('ib', 'deribit')
    }
    flumes: dict[str, Flume] = {
        fqme: Flume(
            mkt=SimpleNamespace(broker=broker),
            first_quote={},
            _rt_shm_token=SimpleNamespace(),
            stream=roots[broker],
        )
        for fqme, broker in (
            ('nvda.nasdaq.ib', 'ib'),
            ('aapl.nasdaq.ib', 'ib'),
            ('btc-usd.deribit', 'deribit'),
        )
    }
    return (
        feed_mod.Feed(
            mods={
                broker: SimpleNamespace(name=broker)
                for broker in roots
            },
            portals={},
            flumes=flumes,
            streams=roots,
            status={},
        ),
        roots,
    )


@pytest.mark.parametrize('cache_hit', [False, True])
def test_cached_feed_returns_local_streams(
    monkeypatch,
    cache_hit: bool,
) -> None:
    '''
    A cached feed caller must not consume or replace shared streams.

    `maybe_open_feed()` previously paired provider subscriptions with
    `Feed.flumes.values()` using `zip()` and installed each child on the
    cached `Flume`. One caller therefore exposed its lexical receiver
    to every holder of the shared `Feed`; after exit those descriptors
    pointed at closed children. The positional pairing also dropped a
    market whenever one provider supplied multiple flumes.

    Model two real provider broadcasters with two IB markets and one
    Deribit market. Enter both first-owner and cache-hit paths, then
    prove the yielded `Feed`, its `Feed.streams` mapping and all
    `Flume`s are caller-local views. Exercise provider-context-wide
    pause/resume and `Feed.open_multi_stream()` to show duplex sends
    still reach the roots while each relay owns another child. Finally
    prove every lexical child closes and all cached descriptors still
    point at live roots.

    '''
    cached_feed: feed_mod.Feed
    roots: dict[str, FakeMsgStream]
    cached_feed, roots = make_feed()
    cached_flumes: dict[str, Flume] = cached_feed.flumes

    @acm
    async def maybe_open_context(**kwargs):
        '''
        Return the controlled shared feed as a cache hit.

        '''
        yield cache_hit, cached_feed

    monkeypatch.setattr(
        feed_mod.trionics,
        'maybe_open_context',
        maybe_open_context,
    )

    local_feed: feed_mod.Feed|None = None
    local_streams: dict[
        str,
        tractor.trionics.BroadcastReceiver,
    ]|None = None

    async def main() -> None:
        nonlocal local_feed, local_streams
        feed: feed_mod.Feed
        async with feed_mod.maybe_open_feed(
            list(cached_flumes),
        ) as feed:
            local_feed = feed
            local_streams = feed.streams
            assert feed is not cached_feed
            assert feed.streams is not roots
            assert feed.flumes is not cached_flumes

            broker: str
            stream: tractor.trionics.BroadcastReceiver
            for broker, stream in feed.streams.items():
                assert stream is roots[broker].children[0]
                assert stream._ctx is roots[broker].ctx

            fqme: str
            flume: Flume
            for fqme, flume in feed.flumes.items():
                cached: Flume = cached_flumes[fqme]

                stream: tractor.trionics.BroadcastReceiver = (
                    feed.streams[flume.mkt.broker]
                )

                assert flume is not cached
                assert flume.stream is stream
                assert cached.stream is roots[flume.mkt.broker]

            assert (
                feed.flumes['nvda.nasdaq.ib'].stream
                is feed.flumes['aapl.nasdaq.ib'].stream
            )

            await feed.pause()
            await feed.resume()
            assert all(
                root.controls == ['pause', 'resume']
                for root in roots.values()
            )

            stream: trio.MemoryReceiveChannel
            async with feed.open_multi_stream() as stream:
                assert all(
                    len(child._state.subs) == 3
                    for child in feed.streams.values()
                )
                await roots['ib'].push({'provider': 'ib'})
                await roots['deribit'].push({
                    'provider': 'deribit',
                })
                with trio.fail_after(1):
                    msgs = {
                        (await stream.receive())['provider'],
                        (await stream.receive())['provider'],
                    }
                assert msgs == {'ib', 'deribit'}

                # One provider's EOC must close only its `tx` clone;
                # the sibling relay continues publishing to `stream`.
                await roots['ib'].close_source()
                with trio.fail_after(1):
                    while (
                        len(feed.streams['ib']._state.subs)
                        != 2
                    ):
                        await trio.lowlevel.checkpoint()

                    await roots['deribit'].push({
                        'provider': 'deribit-after-ib-eoc',
                    })
                    msg: dict
                    msg = await stream.receive()

                assert msg['provider'] == 'deribit-after-ib-eoc'

            assert all(
                len(child._state.subs) == 2
                for child in feed.streams.values()
            )

    trio.run(main)

    assert local_feed is not None
    assert local_streams is not None
    assert all(stream._closed for stream in local_streams.values())
    assert all(
        flume.stream is roots[flume.mkt.broker]
        for flume in cached_flumes.values()
    )
    assert all(
        len(root._broadcaster._state.subs) == 1
        for root in roots.values()
    )


def test_feed_cache_key_includes_shape_and_policy(
    monkeypatch,
) -> None:
    '''
    Feed cache identity must include every resource-defining input.

    The old `fqmes[0]` key aliased `[A]` with `[A, B]` and ignored
    `tick_throttle`, `allow_overruns` and `start_stream`. Whichever
    caller entered first silently determined later feed shape and
    stream behavior, including reuse of a history-only remote feed for
    a live quote request.

    Capture the key passed to `maybe_open_context()` for reordered and
    varied requests. Symbol order must normalize to one key, while the
    full symbol set and each stream policy must produce distinct keys.

    '''
    feed: feed_mod.Feed = feed_mod.Feed(
        mods={},
        portals={},
        flumes={},
        streams={},
        status={},
    )
    keys: list[tuple] = []

    @acm
    async def maybe_open_context(**kwargs):
        '''
        Capture cache identity without opening a remote feed.

        '''
        keys.append(kwargs['key'])
        yield False, feed

    @acm
    async def gather_contexts(mngrs):
        '''
        Permit the metadata-only empty feed used by this key probe.

        '''
        assert not mngrs
        yield ()

    monkeypatch.setattr(
        feed_mod.trionics,
        'maybe_open_context',
        maybe_open_context,
    )
    monkeypatch.setattr(
        feed_mod.trionics,
        'gather_contexts',
        gather_contexts,
    )

    async def open_once(
        fqmes: list[str],
        **kwargs,
    ) -> None:
        async with feed_mod.maybe_open_feed(
            fqmes,
            **kwargs,
        ):
            pass

    async def main() -> None:
        await open_once(['a.ib', 'b.ib'])
        await open_once(['b.ib', 'a.ib'])
        await open_once(['a.ib'])
        await open_once(['a.ib', 'b.ib'], tick_throttle=10)
        await open_once(
            ['a.ib', 'b.ib'],
            allow_overruns=False,
        )
        await open_once(
            ['a.ib', 'b.ib'],
            start_stream=False,
        )

    trio.run(main)

    assert keys[0] == keys[1]
    assert len(set(keys[2:])) == 4
    assert keys[0] not in keys[2:]