test(events): make the timing of the event manager tests deterministic

Author: vdusek

tests/unit/events/test_event_manager.py

@@ -12,15 +12,18 @@
 
 import pytest
 
+from crawlee._utils.wait import wait_for_all_tasks_to_finish
 from crawlee.events import Event, EventManager, EventSystemInfoData
 
 if TYPE_CHECKING:
-    from collections.abc import AsyncGenerator
+    from collections.abc import AsyncGenerator, Sequence
 
 
 @pytest.fixture
 async def event_manager() -> AsyncGenerator[EventManager, None]:
-    async with EventManager() as event_manager:
+    # The teardown waits for the listeners left running by the test, so cap it - without a timeout, a listener that
+    # never completes hangs the whole run instead of failing the test that left it behind.
+    async with EventManager(close_timeout=timedelta(seconds=5)) as event_manager:
         yield event_manager
 
 
@@ -49,6 +52,34 @@ def sync_listener(payload: Any) -> None:
     return sl
 
 
+class ListenerWaitSpy:
+    """Observes the waits the event manager makes, patched in over the real waiting helper on construction.
+
+    It lets a test see which tasks a wait awaits and act at the exact moment a wait starts, so that ordering does
+    not have to be approximated by sleeping.
+    """
+
+    def __init__(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        self.awaited_tasks: list[set[asyncio.Task[None]]] = []
+        """Tasks awaited by each of the waits, in the order the waits started."""
+
+        self.wait_started = asyncio.Event()
+        """Set by a wait right before it starts awaiting, so once it is observed, the waiter is already blocked."""
+
+        monkeypatch.setattr('crawlee.events._event_manager.wait_for_all_tasks_to_finish', self.wait)
+
+    async def wait(
+        self,
+        tasks: Sequence[asyncio.Task[None]],
+        *,
+        logger: logging.Logger,
+        timeout: timedelta | None = None,
+    ) -> None:
+        self.awaited_tasks.append(set(tasks))
+        self.wait_started.set()
+        await wait_for_all_tasks_to_finish(tasks, logger=logger, timeout=timeout)
+
+
 async def test_emit_invokes_registered_sync_listener(
     sync_listener: MagicMock,
     event_manager: EventManager,
@@ -329,24 +360,36 @@ async def raise_error() -> None:
 
 async def test_wait_for_all_listeners_cancelled_error(
     monkeypatch: pytest.MonkeyPatch,
-    caplog: pytest.LogCaptureFixture,
     event_system_info_data: EventSystemInfoData,
 ) -> None:
-    # Simulate long-running listener tasks
-    async def long_running_listener() -> None:
-        await asyncio.sleep(10)
-
-    # Define a side effect function that raises CancelledError
-    async def mock_async_wait(*_: Any, **__: Any) -> None:
+    """A cancelled wait must cancel the listeners it was awaiting and let the cancellation propagate."""
+    never_set = asyncio.Event()
+    listener_cancelled = asyncio.Event()
+
+    async def never_ending_listener() -> None:
+        try:
+            await never_set.wait()
+        except asyncio.CancelledError:
+            listener_cancelled.set()
+            raise
+
+    async def cancel_the_wait(*_: Any, **__: Any) -> None:
         raise asyncio.CancelledError
 
-    with pytest.raises(asyncio.CancelledError), caplog.at_level(logging.WARNING):  # noqa: PT012
-        async with EventManager(close_timeout=timedelta(milliseconds=10)) as event_manager:
-            event_manager.on(event=Event.SYSTEM_INFO, listener=long_running_listener)
-            event_manager.emit(event=Event.SYSTEM_INFO, event_data=event_system_info_data)
+    event_manager = EventManager()
+    await event_manager.__aenter__()
+    event_manager.on(event=Event.SYSTEM_INFO, listener=never_ending_listener)
+    event_manager.emit(event=Event.SYSTEM_INFO, event_data=event_system_info_data)
+
+    # Only `asyncio.wait` is replaced, so the cancellation is still handled by the real waiting code - which is what
+    # has to cancel the listener that never finishes on its own.
+    monkeypatch.setattr('asyncio.wait', cancel_the_wait)
+
+    # Capped, as a listener left uncancelled would be awaited forever - the close has no timeout of its own here.
+    with pytest.raises(asyncio.CancelledError):
+        await asyncio.wait_for(event_manager.__aexit__(None, None, None), timeout=5)
 
-            # Use monkeypatch to replace asyncio.wait with mock_async_wait
-            monkeypatch.setattr('asyncio.wait', mock_async_wait)
+    assert listener_cancelled.is_set()
 
 
 async def test_methods_raise_error_when_not_active(event_system_info_data: EventSystemInfoData) -> None:
@@ -375,16 +418,20 @@ async def test_wait_for_all_listeners_from_within_a_listener_does_not_deadlock(
     event_system_info_data: EventSystemInfoData,
 ) -> None:
     """Waiting from within a listener must not self-await, yet must still await the other listeners."""
+    parked = asyncio.Event()
+    release_other_listener = asyncio.Event()
     other_listener_done = asyncio.Event()
     waiter_done = asyncio.Event()
     other_done_when_wait_returned: bool | None = None
 
     async def other_listener(_: Any) -> None:
-        await asyncio.sleep(0.2)
+        await release_other_listener.wait()
         other_listener_done.set()
 
     async def waiting_listener(_: Any) -> None:
         nonlocal other_done_when_wait_returned
+        # Set right before the wait, which registers the waiter and captures the tasks without yielding in between.
+        parked.set()
         await event_manager.wait_for_all_listeners_to_complete()
         other_done_when_wait_returned = other_listener_done.is_set()
         waiter_done.set()
@@ -393,11 +440,14 @@ async def waiting_listener(_: Any) -> None:
     event_manager.on(event=Event.SYSTEM_INFO, listener=waiting_listener)
     event_manager.emit(event=Event.SYSTEM_INFO, event_data=event_system_info_data)
 
+    # The other listener finishes only once the waiter is already blocked, so the wait cannot pass it by.
+    await asyncio.wait_for(parked.wait(), timeout=5)
+    release_other_listener.set()
+
     await asyncio.wait_for(waiter_done.wait(), timeout=5)
 
     # No self-await deadlock, and the wait must have blocked until the co-registered listener finished.
     assert other_done_when_wait_returned is True
-    assert other_listener_done.is_set()
 
 
 async def test_wait_from_within_multiple_listeners_does_not_deadlock(
@@ -427,39 +477,50 @@ async def second_waiting_listener(_: Any) -> None:
 
 
 async def test_wait_from_outside_awaits_a_listener_that_is_itself_waiting(
+    monkeypatch: pytest.MonkeyPatch,
     event_manager: EventManager,
     event_system_info_data: EventSystemInfoData,
 ) -> None:
     """A caller that is not a listener is outside the deadlock cycle, so it must await even waiting listeners."""
+    spy = ListenerWaitSpy(monkeypatch)
     parked = asyncio.Event()
+    release_other_listener = asyncio.Event()
     waiting_listener_done = asyncio.Event()
+    waiter_task: asyncio.Task[None] | None = None
 
     async def other_listener(_: Any) -> None:
-        await asyncio.sleep(0.1)
+        await release_other_listener.wait()
 
     async def waiting_listener(_: Any) -> None:
+        nonlocal waiter_task
+        waiter_task = asyncio.current_task()
+        # Set right before the wait, which registers the waiter and captures the tasks without yielding in between.
         parked.set()
         await event_manager.wait_for_all_listeners_to_complete()
-        # Work done after the inner wait returns - the outer wait must not return before it finishes.
-        await asyncio.sleep(0.1)
         waiting_listener_done.set()
 
     event_manager.on(event=Event.SYSTEM_INFO, listener=other_listener)
     event_manager.on(event=Event.SYSTEM_INFO, listener=waiting_listener)
     event_manager.emit(event=Event.SYSTEM_INFO, event_data=event_system_info_data)
 
-    # `parked` is set right before the listener registers itself as a waiter, without yielding in between.
     await asyncio.wait_for(parked.wait(), timeout=5)
 
+    # The release has to come from outside this task, as the wait below blocks until the listeners are done. Scheduled
+    # on the loop, it runs only once that wait has suspended - so the wait captures the tasks with both still running.
+    asyncio.get_running_loop().call_soon(release_other_listener.set)
     await asyncio.wait_for(event_manager.wait_for_all_listeners_to_complete(), timeout=5)
 
+    # The waiter's own wait went first, so the second one is the outer wait - and it had to include the parked waiter.
+    assert waiter_task in spy.awaited_tasks[1]
     assert waiting_listener_done.is_set()
 
 
 async def test_close_from_within_a_listener_does_not_deadlock_or_error(
+    monkeypatch: pytest.MonkeyPatch,
     event_system_info_data: EventSystemInfoData,
 ) -> None:
     """Closing the event manager from within a listener (as `Actor.exit()` does) must not deadlock or raise."""
+    spy = ListenerWaitSpy(monkeypatch)
     event_manager = EventManager()
     await event_manager.__aenter__()
 
@@ -468,30 +529,35 @@ async def test_close_from_within_a_listener_does_not_deadlock_or_error(
     asyncio.get_running_loop().set_exception_handler(lambda _loop, context: loop_errors.append(context))
 
     closed = asyncio.Event()
+    release_other_listener = asyncio.Event()
     other_listener_done = asyncio.Event()
+    closing_task: asyncio.Task[None] | None = None
 
     async def other_listener(_: Any) -> None:
-        await asyncio.sleep(0.2)
+        await release_other_listener.wait()
         other_listener_done.set()
 
     async def closing_listener(_: Any) -> None:
+        nonlocal closing_task
+        closing_task = asyncio.current_task()
         await event_manager.__aexit__(None, None, None)
         closed.set()
 
     # A second listener makes close await a concurrently-running listener - the real `Actor.exit()` shape.
     event_manager.on(event=Event.SYSTEM_INFO, listener=other_listener)
     event_manager.on(event=Event.SYSTEM_INFO, listener=closing_listener)
-
-    tasks_before = asyncio.all_tasks()
     event_manager.emit(event=Event.SYSTEM_INFO, event_data=event_system_info_data)
 
     try:
+        # The close is already blocked in its wait here, so releasing the other listener only now means the close
+        # really does have to await a listener that is still running.
+        await asyncio.wait_for(spy.wait_started.wait(), timeout=5)
+        release_other_listener.set()
+
         await asyncio.wait_for(closed.wait(), timeout=5)
-        # Drain the wrapper tasks so their `finally` blocks run before we assert - no arbitrary sleep.
-        spawned = asyncio.all_tasks() - tasks_before - {asyncio.current_task()}
-        if spawned:
-            # Capped, as the set is derived from all the running tasks - an unrelated one must not hang the test.
-            await asyncio.wait(spawned, timeout=5)
+        # Let the closing listener's own task finalize before asserting - its wrapper returns right after `closed`.
+        assert closing_task is not None
+        await asyncio.wait_for(closing_task, timeout=5)
     finally:
         # Cap the cleanup so a regressed deadlock surfaces the real failure instead of hanging.
         if event_manager.active:

tests/unit/events/test_local_event_manager.py

@@ -14,21 +14,25 @@
     import pytest
 
 
-async def test_emit_system_info_event() -> None:
+async def test_emit_system_info_event(monkeypatch: pytest.MonkeyPatch) -> None:
+    """The recurring task emits the first `SystemInfo` event as soon as it starts, without waiting for the interval."""
+    # Both readings are replaced with instant ones - a real `get_cpu_info` samples the CPU utilization over 100 ms,
+    # and on a loaded runner it takes far longer than that.
+    monkeypatch.setattr('crawlee.events._local_event_manager.get_cpu_info', lambda: MagicMock(spec=CpuInfo))
+    monkeypatch.setattr('crawlee.events._local_event_manager.get_memory_info', lambda: MagicMock(spec=MemoryInfo))
+
     mocked_listener = AsyncMock()
     received = asyncio.Event()
 
     async def async_listener(payload: Any) -> None:
         await mocked_listener(payload)
         received.set()
 
-    system_info_interval = timedelta(milliseconds=50)
-    async with LocalEventManager(system_info_interval=system_info_interval) as event_manager:
+    # An interval this long means the event can only come from the immediate first run of the recurring task.
+    async with LocalEventManager(system_info_interval=timedelta(hours=1)) as event_manager:
+        # Registered before anything yields to the event loop, so the very first emission already reaches it.
         event_manager.on(event=Event.SYSTEM_INFO, listener=async_listener)
-        # Wait until the listener is invoked at least once. A generous timeout avoids flakiness on
-        # loaded CI runners, where `psutil.cpu_percent(interval=0.1)` in `asyncio.to_thread` can
-        # delay the first emission well beyond the configured interval.
-        await asyncio.wait_for(received.wait(), timeout=10)
+        await asyncio.wait_for(received.wait(), timeout=5)
 
     assert mocked_listener.call_count >= 1
     assert isinstance(mocked_listener.call_args[0][0], EventSystemInfoData)