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.

tests/test_local_task_broadcast.py

@@ -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(
+@asynccontextmanager
-    arb_addr,
+async def open_sequence_streamer(
-    start_method,
-):
 
-    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():