Update test_broker.py

Author: sobolevn

tests/abc/test_broker.py

@@ -83,23 +83,31 @@ 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()
     broker.is_worker_process = is_worker_process
     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,23 +121,31 @@ 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()
     broker.is_worker_process = is_worker_process
     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