from __future__ import annotations import asyncio import logging from datetime import timedelta from functools import update_wrapper from typing import TYPE_CHECKING, Any from unittest import mock from unittest.mock import AsyncMock, MagicMock import pytest from crawlee.events import Event, EventManager, EventSystemInfoData if TYPE_CHECKING: from collections.abc import AsyncGenerator @pytest.fixture async def event_manager() -> AsyncGenerator[EventManager, None]: async with EventManager() as event_manager: yield event_manager @pytest.fixture def event_system_info_data() -> EventSystemInfoData: return MagicMock(spec=EventSystemInfoData) @pytest.fixture def async_listener() -> AsyncMock: async def async_listener(payload: Any) -> None: pass al = AsyncMock() update_wrapper(al, async_listener) return al @pytest.fixture def sync_listener() -> MagicMock: def sync_listener(payload: Any) -> None: pass sl = MagicMock() update_wrapper(sl, sync_listener) return sl async def test_emit_invokes_registered_sync_listener( sync_listener: MagicMock, event_manager: EventManager, event_system_info_data: EventSystemInfoData, ) -> None: event_manager.on(event=Event.SYSTEM_INFO, listener=sync_listener) event_manager.emit(event=Event.SYSTEM_INFO, event_data=event_system_info_data) await event_manager.wait_for_all_listeners_to_complete() assert sync_listener.call_count == 1 assert sync_listener.call_args[0] == (event_system_info_data,) async def test_emit_invokes_both_sync_and_async_listeners( sync_listener: MagicMock, async_listener: AsyncMock, event_manager: EventManager, event_system_info_data: EventSystemInfoData, ) -> None: event_manager.on(event=Event.SYSTEM_INFO, listener=sync_listener) event_manager.on(event=Event.SYSTEM_INFO, listener=async_listener) event_manager.emit(event=Event.SYSTEM_INFO, event_data=event_system_info_data) await event_manager.wait_for_all_listeners_to_complete() assert async_listener.call_count == 1 assert async_listener.call_args[0] == (event_system_info_data,) assert sync_listener.call_count == 1 assert sync_listener.call_args[0] == (event_system_info_data,) async def test_emit_event_with_no_listeners( event_manager: EventManager, event_system_info_data: EventSystemInfoData, async_listener: AsyncMock, ) -> None: # Register a listener for a different event event_manager.on(event=Event.ABORTING, listener=async_listener) # Attempt to emit an event for which no listeners are registered, it should not fail event_manager.emit(event=Event.SYSTEM_INFO, event_data=event_system_info_data) await event_manager.wait_for_all_listeners_to_complete() # Ensure the listener for the other event was not called assert async_listener.call_count == 0 async def test_emit_invokes_parameterless_listener( event_manager: EventManager, event_system_info_data: EventSystemInfoData, ) -> None: sync_mock = MagicMock() def sync_listener() -> None: sync_mock() async_mock = MagicMock() async def async_listener() -> None: async_mock() event_manager.on(event=Event.SYSTEM_INFO, listener=sync_listener) event_manager.on(event=Event.SYSTEM_INFO, listener=async_listener) event_manager.emit(event=Event.SYSTEM_INFO, event_data=event_system_info_data) await event_manager.wait_for_all_listeners_to_complete() assert sync_mock.call_count == 1 assert async_mock.call_count == 1 async def test_remove_nonexistent_listener_does_not_fail( async_listener: AsyncMock, event_manager: EventManager, ) -> None: # Attempt to remove a specific listener that was never added. event_manager.off(event=Event.SYSTEM_INFO, listener=async_listener) # Attempt to remove all listeners. event_manager.off(event=Event.ABORTING) async def test_removed_listener_not_invoked_on_emit( async_listener: AsyncMock, event_manager: EventManager, event_system_info_data: EventSystemInfoData, ) -> None: event_manager.on(event=Event.SYSTEM_INFO, listener=async_listener) event_manager.off(event=Event.SYSTEM_INFO, listener=async_listener) event_manager.emit(event=Event.SYSTEM_INFO, event_data=event_system_info_data) await event_manager.wait_for_all_listeners_to_complete() assert async_listener.call_count == 0 async def test_close_clears_listeners_and_tasks(async_listener: AsyncMock) -> None: async with EventManager() as event_manager: event_manager.on(event=Event.SYSTEM_INFO, listener=async_listener) assert async_listener.call_count == 0 assert len(event_manager._listener_tasks) == 0 assert len(event_manager._listeners_to_wrappers) == 0 async def test_close_after_emit_processes_event( async_listener: AsyncMock, event_system_info_data: EventSystemInfoData, ) -> None: async with EventManager() as event_manager: event_manager.on(event=Event.SYSTEM_INFO, listener=async_listener) event_manager.emit(event=Event.SYSTEM_INFO, event_data=event_system_info_data) # Event should be processed before the event manager is closed assert async_listener.call_count == 1 assert async_listener.call_args[0] == (event_system_info_data,) assert len(event_manager._listener_tasks) == 0 assert len(event_manager._listeners_to_wrappers) == 0 async def test_wait_for_all_listeners_cancelled_error( monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, ) -> 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: 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) # Use monkeypatch to replace asyncio.wait with mock_async_wait monkeypatch.setattr('asyncio.wait', mock_async_wait) async def test_methods_raise_error_when_not_active(event_system_info_data: EventSystemInfoData) -> None: event_manager = EventManager() assert event_manager.active is False with pytest.raises(RuntimeError, match=r'EventManager is not active.'): event_manager.emit(event=Event.SYSTEM_INFO, event_data=event_system_info_data) with pytest.raises(RuntimeError, match=r'EventManager is not active.'): await event_manager.wait_for_all_listeners_to_complete() async with event_manager: event_manager.emit(event=Event.SYSTEM_INFO, event_data=event_system_info_data) await event_manager.wait_for_all_listeners_to_complete() assert event_manager.active is True async def test_event_manager_in_context_persistence() -> None: """Test that entering the `EventManager` context emits persist state event at least once.""" event_manager = EventManager() with mock.patch.object(event_manager, '_emit_persist_state_event', AsyncMock()) as mocked_emit_persist_state_event: async with event_manager: pass assert mocked_emit_persist_state_event.call_count >= 1