From 0a214883efcb3ec2bc3e6097ddd950f72eaf49cb Mon Sep 17 00:00:00 2001 From: sobolevn Date: Mon, 4 May 2026 15:50:01 +0300 Subject: [PATCH 1/3] Refactor tests for `async with broker` Refs https://github.com/taskiq-python/taskiq/commit/653400e7d440c850098747a873efb6ad3537cabd --- tests/abc/test_broker.py | 88 +++++++++++----------------------------- 1 file changed, 23 insertions(+), 65 deletions(-) diff --git a/tests/abc/test_broker.py b/tests/abc/test_broker.py index 8d39ea50..c6818ef2 100644 --- a/tests/abc/test_broker.py +++ b/tests/abc/test_broker.py @@ -83,77 +83,52 @@ async def test_task() -> None: ... @pytest.mark.anyio -async def test_async_context_manager_enter() -> None: - """Test that __aenter__ calls startup.""" +@pytest.mark.parametrize('is_worker_process', [True, False]) +async def test_async_context_manager_enter( + *, + is_worker_process: bool, +) -> None: + """Test that `__aenter__` and `__aexit__` calls work.""" broker = _TestBroker() + broker.is_worker_process = is_worker_process startup_called = False + shutdown_called = False @broker.on_event(TaskiqEvents.CLIENT_STARTUP) async def track_startup(state: TaskiqState) -> None: nonlocal startup_called startup_called = True - async with broker: - assert startup_called is True - - -@pytest.mark.anyio -async def test_async_context_manager_exit() -> None: - """Test that __aexit__ calls shutdown.""" - broker = _TestBroker() - shutdown_called = False - @broker.on_event(TaskiqEvents.CLIENT_SHUTDOWN) async def track_shutdown(state: TaskiqState) -> None: nonlocal shutdown_called shutdown_called = True - async with broker: - pass + async with broker as ctx: + assert ctx is None + assert startup_called is True + assert shutdown_called is False assert shutdown_called is True @pytest.mark.anyio -async def test_async_context_manager_enter_worker() -> None: - """Test that __aenter__ calls worker startup when is_worker_process is True.""" +@pytest.mark.parametrize('is_worker_process', [True, False]) +async def test_async_context_manager_exit_on_exception( + *, + is_worker_process: bool, +) -> None: + """Test that __aexit__ calls shutdown even if exception is raised.""" broker = _TestBroker() - broker.is_worker_process = True + broker.is_worker_process = is_worker_process startup_called = False + shutdown_called = False - @broker.on_event(TaskiqEvents.WORKER_STARTUP) + @broker.on_event(TaskiqEvents.CLIENT_STARTUP) async def track_startup(state: TaskiqState) -> None: nonlocal startup_called startup_called = True - async with broker: - assert startup_called is True - - -@pytest.mark.anyio -async def test_async_context_manager_exit_worker() -> None: - """Test that __aexit__ calls worker shutdown when is_worker_process is True.""" - broker = _TestBroker() - broker.is_worker_process = True - shutdown_called = False - - @broker.on_event(TaskiqEvents.WORKER_SHUTDOWN) - async def track_shutdown(state: TaskiqState) -> None: - nonlocal shutdown_called - shutdown_called = True - - async with broker: - pass - - assert shutdown_called is True - - -@pytest.mark.anyio -async def test_async_context_manager_exit_on_exception() -> None: - """Test that __aexit__ calls shutdown even if exception is raised.""" - broker = _TestBroker() - shutdown_called = False - @broker.on_event(TaskiqEvents.CLIENT_SHUTDOWN) async def track_shutdown(state: TaskiqState) -> None: nonlocal shutdown_called @@ -161,25 +136,8 @@ async def track_shutdown(state: TaskiqState) -> None: with pytest.raises(ValueError, match="Test exception"): async with broker: - raise ValueError("Test exception") - - assert shutdown_called is True - - -@pytest.mark.anyio -async def test_async_context_manager_exit_worker_on_exception() -> None: - """Test that __aexit__ calls worker shutdown even if exception is raised.""" - broker = _TestBroker() - broker.is_worker_process = True - shutdown_called = False - - @broker.on_event(TaskiqEvents.WORKER_SHUTDOWN) - async def track_shutdown(state: TaskiqState) -> None: - nonlocal shutdown_called - shutdown_called = True - - with pytest.raises(ValueError, match="Test exception"): - async with broker: + assert startup_called is True + assert shutdown_called is False raise ValueError("Test exception") assert shutdown_called is True From 58d2e6bd0a221445863d4692d15c612cacadcd4a Mon Sep 17 00:00:00 2001 From: sobolevn Date: Mon, 4 May 2026 15:54:05 +0300 Subject: [PATCH 2/3] Update test_broker.py --- tests/abc/test_broker.py | 28 ++++++++++++++++++++++------ 1 file changed, 22 insertions(+), 6 deletions(-) diff --git a/tests/abc/test_broker.py b/tests/abc/test_broker.py index c6818ef2..16526bd2 100644 --- a/tests/abc/test_broker.py +++ b/tests/abc/test_broker.py @@ -83,10 +83,18 @@ async def test_task() -> None: ... @pytest.mark.anyio -@pytest.mark.parametrize('is_worker_process', [True, False]) +@pytest.mark.parametrize( + ('is_worker_process', 'startup', 'shutdown'), + [ + (True, TaskiqEvents.WORKER_STARTUP, TaskiqEvents.WORKER_SHUTDOWN), + (False, TaskiqEvents.CLIENT_STARTUP, TaskiqEvents.CLIENT_SHUTDOWN), + ], +) async def test_async_context_manager_enter( *, is_worker_process: bool, + startup: TaskiqEvents, + shutdown: TaskiqEvents, ) -> None: """Test that `__aenter__` and `__aexit__` calls work.""" broker = _TestBroker() @@ -94,12 +102,12 @@ async def test_async_context_manager_enter( startup_called = False shutdown_called = False - @broker.on_event(TaskiqEvents.CLIENT_STARTUP) + @broker.on_event(startup) async def track_startup(state: TaskiqState) -> None: nonlocal startup_called startup_called = True - @broker.on_event(TaskiqEvents.CLIENT_SHUTDOWN) + @broker.on_event(shutdown) async def track_shutdown(state: TaskiqState) -> None: nonlocal shutdown_called shutdown_called = True @@ -113,10 +121,18 @@ async def track_shutdown(state: TaskiqState) -> None: @pytest.mark.anyio -@pytest.mark.parametrize('is_worker_process', [True, False]) +@pytest.mark.parametrize( + ('is_worker_process', 'startup', 'shutdown'), + [ + (True, TaskiqEvents.WORKER_STARTUP, TaskiqEvents.WORKER_SHUTDOWN), + (False, TaskiqEvents.CLIENT_STARTUP, TaskiqEvents.CLIENT_SHUTDOWN), + ], +) async def test_async_context_manager_exit_on_exception( *, is_worker_process: bool, + startup: TaskiqEvents, + shutdown: TaskiqEvents, ) -> None: """Test that __aexit__ calls shutdown even if exception is raised.""" broker = _TestBroker() @@ -124,12 +140,12 @@ async def test_async_context_manager_exit_on_exception( startup_called = False shutdown_called = False - @broker.on_event(TaskiqEvents.CLIENT_STARTUP) + @broker.on_event(startup) async def track_startup(state: TaskiqState) -> None: nonlocal startup_called startup_called = True - @broker.on_event(TaskiqEvents.CLIENT_SHUTDOWN) + @broker.on_event(shutdown) async def track_shutdown(state: TaskiqState) -> None: nonlocal shutdown_called shutdown_called = True From 4d49422f39be3ac7990025e7e8c8c19f59927768 Mon Sep 17 00:00:00 2001 From: sobolevn Date: Mon, 4 May 2026 16:13:22 +0300 Subject: [PATCH 3/3] Fix ruff --- tests/abc/test_broker.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/abc/test_broker.py b/tests/abc/test_broker.py index 16526bd2..636f9576 100644 --- a/tests/abc/test_broker.py +++ b/tests/abc/test_broker.py @@ -84,7 +84,7 @@ async def test_task() -> None: ... @pytest.mark.anyio @pytest.mark.parametrize( - ('is_worker_process', 'startup', 'shutdown'), + ("is_worker_process", "startup", "shutdown"), [ (True, TaskiqEvents.WORKER_STARTUP, TaskiqEvents.WORKER_SHUTDOWN), (False, TaskiqEvents.CLIENT_STARTUP, TaskiqEvents.CLIENT_SHUTDOWN), @@ -122,7 +122,7 @@ async def track_shutdown(state: TaskiqState) -> None: @pytest.mark.anyio @pytest.mark.parametrize( - ('is_worker_process', 'startup', 'shutdown'), + ("is_worker_process", "startup", "shutdown"), [ (True, TaskiqEvents.WORKER_STARTUP, TaskiqEvents.WORKER_SHUTDOWN), (False, TaskiqEvents.CLIENT_STARTUP, TaskiqEvents.CLIENT_SHUTDOWN),