forked from goodboy/tractor
1
0
Fork 0

Add laggy parent stream tests

Add a couple more tests to check that a parent and sub-task stream can
be lagged and recovered (depending on who's slower). Factor some of the
test machinery into a new ctx mngr to make it all happen.
live_on_air_from_tokio
Tyler Goodlet 2021-08-31 13:03:20 -04:00
parent 093e7d921c
commit 0d70e3081a
1 changed files with 145 additions and 34 deletions

View File

@ -1,10 +1,13 @@
"""
Broadcast channels for fan-out to local tasks.
"""
from contextlib import asynccontextmanager
from functools import partial
from itertools import cycle
import time
from typing import Optional
import pytest
import trio
from trio.lowlevel import current_task
import tractor
@ -32,8 +35,11 @@ async def echo_sequences(
async def ensure_sequence(
stream: tractor.ReceiveMsgStream,
sequence: list,
delay: Optional[float] = None,
) -> None:
name = current_task().name
@ -44,19 +50,22 @@ async def ensure_sequence(
assert value == sequence[0]
sequence.remove(value)
if delay:
await trio.sleep(delay)
if not sequence:
# fully consumed
break
def test_stream_fan_out_to_local_subscriptions(
arb_addr,
start_method,
):
@asynccontextmanager
async def open_sequence_streamer(
sequence = list(range(1000))
sequence: list[int],
arb_addr: tuple[str, int],
start_method: str,
async def main():
) -> tractor.MsgStream:
async with tractor.open_nursery(
arbiter_addr=arb_addr,
@ -74,6 +83,25 @@ def test_stream_fan_out_to_local_subscriptions(
assert first is None
async with ctx.open_stream() as stream:
yield stream
await portal.cancel_actor()
def test_stream_fan_out_to_local_subscriptions(
arb_addr,
start_method,
):
sequence = list(range(1000))
async def main():
async with open_sequence_streamer(
sequence,
arb_addr,
start_method,
) as stream:
async with trio.open_nursery() as n:
for i in range(10):
@ -95,11 +123,94 @@ def test_stream_fan_out_to_local_subscriptions(
# fully consumed
break
await portal.cancel_actor()
trio.run(main)
@pytest.mark.parametrize(
'task_delays',
[
(0.01, 0.001),
(0.001, 0.01),
]
)
def test_consumer_and_parent_maybe_lag(
arb_addr,
start_method,
task_delays,
):
async def main():
sequence = list(range(1000))
parent_delay, sub_delay = task_delays
async with open_sequence_streamer(
sequence,
arb_addr,
start_method,
) as stream:
try:
async with trio.open_nursery() as n:
n.start_soon(
ensure_sequence,
stream,
sequence.copy(),
sub_delay,
name='consumer_task',
)
await stream.send(tuple(sequence))
# async for value in stream:
lagged = False
lag_count = 0
while True:
try:
value = await stream.receive()
print(f'source stream rx: {value}')
if lagged:
# re set the sequence starting at our last
# value
sequence = sequence[sequence.index(value) + 1:]
else:
assert value == sequence[0]
sequence.remove(value)
lagged = False
except Lagged:
lagged = True
print(f'source stream lagged after {value}')
lag_count += 1
continue
# lag the parent
await trio.sleep(parent_delay)
if not sequence:
# fully consumed
break
print(f'parent + source stream lagged: {lag_count}')
if parent_delay > sub_delay:
assert lag_count > 0
except Lagged:
# child was lagged
assert parent_delay < sub_delay
trio.run(main)
# TODO:
# def test_first_task_to_recv_is_cancelled():
# ...
def test_subscribe_errors_after_close():
async def main():