121 lines
3.6 KiB
Python
121 lines
3.6 KiB
Python
'''
|
|
IB cached quote-channel ownership regressions.
|
|
|
|
'''
|
|
from contextlib import asynccontextmanager as acm
|
|
|
|
import tractor
|
|
import trio
|
|
|
|
from piker.brokers.ib import feed
|
|
|
|
|
|
def test_quote_stream_caches_channel_and_owns_child(
|
|
monkeypatch,
|
|
) -> None:
|
|
'''
|
|
IB quote callers must share one cached channel through child cursors.
|
|
|
|
`open_aio_quote_stream()` previously stored the first caller's raw
|
|
`LinkedTaskChannel` in `_quote_streams`, yielded that root directly,
|
|
and wrapped the same source in a new broadcaster for each later
|
|
caller. The first caller therefore owned source lifetime while
|
|
competing roots split quote receives between concurrent users.
|
|
|
|
Stub `maybe_open_context()` with one retained channel and enter a
|
|
miss owner plus concurrent hit. Prove both receive the same quote
|
|
through distinct children of one broadcaster, the hit closes
|
|
independently, and the owner keeps receiving until its own exit.
|
|
Also verify both calls request the same symbol cache key.
|
|
|
|
'''
|
|
class FakeLinkedTaskChannel:
|
|
'''
|
|
Retain one real broadcaster and expose child subscriptions.
|
|
|
|
'''
|
|
def __init__(self) -> None:
|
|
self._tx: trio.MemorySendChannel
|
|
rx: trio.MemoryReceiveChannel
|
|
self._tx, rx = trio.open_memory_channel(8)
|
|
self._broadcaster: (
|
|
tractor.trionics.BroadcastReceiver
|
|
) = tractor.trionics.broadcast_receiver(
|
|
rx,
|
|
8,
|
|
)
|
|
|
|
@acm
|
|
async def subscribe(self):
|
|
'''
|
|
Yield one caller-owned child from the retained root.
|
|
|
|
'''
|
|
child: tractor.trionics.BroadcastReceiver
|
|
async with self._broadcaster.subscribe() as child:
|
|
yield child
|
|
|
|
async def push(self, msg: dict) -> None:
|
|
'''
|
|
Send one ticker through the retained source channel.
|
|
|
|
'''
|
|
await self._tx.send(msg)
|
|
|
|
chan: FakeLinkedTaskChannel = FakeLinkedTaskChannel()
|
|
context_entries: int = 0
|
|
keys: list[str] = []
|
|
acm_funcs: list = []
|
|
|
|
@acm
|
|
async def maybe_open_context(**kwargs):
|
|
'''
|
|
Return one retained channel as a miss followed by a hit.
|
|
|
|
'''
|
|
nonlocal context_entries
|
|
keys.append(kwargs['key'])
|
|
acm_funcs.append(kwargs['acm_func'])
|
|
cache_hit: bool = context_entries > 0
|
|
context_entries += 1
|
|
yield cache_hit, chan
|
|
|
|
monkeypatch.setattr(
|
|
feed.tractor.trionics,
|
|
'maybe_open_context',
|
|
maybe_open_context,
|
|
)
|
|
|
|
async def main() -> None:
|
|
owner: tractor.trionics.BroadcastReceiver
|
|
async with feed.open_aio_quote_stream(
|
|
'NVDA',
|
|
) as owner:
|
|
hit: tractor.trionics.BroadcastReceiver
|
|
async with feed.open_aio_quote_stream(
|
|
'NVDA',
|
|
) as hit:
|
|
assert hit is not owner
|
|
assert len(chan._broadcaster._state.subs) == 3
|
|
|
|
await chan.push({'value': 1})
|
|
assert await owner.receive() == {'value': 1}
|
|
assert await hit.receive() == {'value': 1}
|
|
|
|
assert hit._closed
|
|
assert not owner._closed
|
|
assert len(chan._broadcaster._state.subs) == 2
|
|
|
|
await chan.push({'value': 2})
|
|
assert await owner.receive() == {'value': 2}
|
|
|
|
assert owner._closed
|
|
assert len(chan._broadcaster._state.subs) == 1
|
|
assert keys == ['NVDA', 'NVDA']
|
|
assert acm_funcs == [
|
|
feed._open_aio_quote_channel,
|
|
feed._open_aio_quote_channel,
|
|
]
|
|
|
|
trio.run(main)
|