piker/tests/test_deribit_broadcast.py

113 lines
3.4 KiB
Python
Raw Normal View History

'''
Deribit cached feed broadcast ownership regressions.
'''
from contextlib import asynccontextmanager as acm
import tractor
import trio
from piker.brokers.deribit import api
def test_price_feed_owns_child_subscription(
monkeypatch,
) -> None:
'''
Every cached price-feed caller must own one lexical child receiver.
`maybe_open_price_feed()` previously yielded the raw
`LinkedTaskChannel` to the first cache owner, then constructed a new
broadcaster root around that same channel for every cache hit.
Concurrent roots competed for source receives, so quotes could be
split across callers and exiting one wrapper did not express child
ownership on the retained channel broadcaster.
Return one real `LinkedTaskChannel`-shaped fake first as a miss and
then as a concurrent hit. Nest both caller contexts and prove each
gets a distinct child from the channel's single retained root and
receives the same quote. Close the hit child, then prove the first
owner continues receiving before its own lexical child closes.
'''
class FakeLinkedTaskChannel:
'''
Retain one real broadcaster and expose its public subscription.
'''
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 quote through the retained source channel.
'''
await self._tx.send(msg)
channel: FakeLinkedTaskChannel = FakeLinkedTaskChannel()
context_entries: int = 0
@acm
async def maybe_open_context(**kwargs):
'''
Return the controlled channel with either cache state.
'''
nonlocal context_entries
cache_hit: bool = context_entries > 0
context_entries += 1
yield cache_hit, channel
monkeypatch.setattr(
api,
'maybe_open_context',
maybe_open_context,
)
async def main() -> None:
owner: tractor.trionics.BroadcastReceiver
async with api.maybe_open_price_feed(
'btc-usd',
) as owner:
hit: tractor.trionics.BroadcastReceiver
async with api.maybe_open_price_feed(
'btc-usd',
) as hit:
assert hit is not owner
assert len(channel._broadcaster._state.subs) == 3
await channel.push({'value': 1})
assert await owner.receive() == {'value': 1}
assert await hit.receive() == {'value': 1}
assert hit._closed
assert not owner._closed
assert len(channel._broadcaster._state.subs) == 2
await channel.push({'value': 2})
assert await owner.receive() == {'value': 2}
assert owner._closed
assert len(channel._broadcaster._state.subs) == 1
trio.run(main)