From dc4bae2b756e2f81cb1407f3f39e2cd5e04bf3ce Mon Sep 17 00:00:00 2001 From: "sergei.romanchuk" Date: Mon, 13 Jul 2026 21:48:55 +0200 Subject: [PATCH 1/2] fix(inmemory): stabilize task ordering and lifecycle Ensure post_send completes before local execution while preserving per-invocation ownership and await_inplace semantics. Drain accepted tasks during shutdown, retain background failures, and clean up middleware, result backend, and executor resources reliably. Add concurrency, cancellation, lifecycle, and compatibility coverage. Closes: #586 --- docs/available-components/brokers.md | 5 + docs/guide/testing-taskiq.md | 16 + taskiq/abc/broker.py | 81 +++- taskiq/brokers/inmemory_broker.py | 241 ++++++++++- taskiq/kicker.py | 24 +- tests/abc/test_broker.py | 94 +++- tests/brokers/inmemory_contract_support.py | 182 ++++++++ tests/brokers/test_inmemory_lifecycle.py | 301 +++++++++++++ tests/brokers/test_inmemory_send.py | 400 ++++++++++++++++++ .../middlewares/admin_middleware/conftest.py | 13 +- .../test_arguments_formatting.py | 1 - tests/middlewares/test_simple_retry.py | 30 +- 12 files changed, 1316 insertions(+), 72 deletions(-) create mode 100644 tests/brokers/inmemory_contract_support.py create mode 100644 tests/brokers/test_inmemory_lifecycle.py create mode 100644 tests/brokers/test_inmemory_send.py diff --git a/docs/available-components/brokers.md b/docs/available-components/brokers.md index c4dad6df..2108efd0 100644 --- a/docs/available-components/brokers.md +++ b/docs/available-components/brokers.md @@ -12,6 +12,11 @@ This is a special broker for local development. It uses the same functions to ex but all tasks are executed locally in the current thread. By default it uses `InMemoryResultBackend` but this can be overridden. +`startup()` and `shutdown()` run both the client and worker event phases because +the broker performs both roles in one process. They also own middleware and +result backend lifecycle. Shutdown waits for local background executions before +closing those resources and the synchronous task executor. + ## ZeroMQBroker This broker uses [ZMQ](https://zeromq.org/) to communicate between worker and client processes. diff --git a/docs/guide/testing-taskiq.md b/docs/guide/testing-taskiq.md index 14ae47fe..de9a4268 100644 --- a/docs/guide/testing-taskiq.md +++ b/docs/guide/testing-taskiq.md @@ -165,6 +165,10 @@ broker = InMemoryBroker(await_inplace=True) With this setup all `await function.kiq()` calls will behave similarly to `await function()`, but with dependency injection and all taskiq-related functionality. +The normal middleware boundary is preserved in both execution modes: all +`post_send` hooks finish before `pre_execute` and the task body begin. +With `await_inplace=True`, the `kiq()` call also returns only after inline +execution completes. 2. Alternatively, you can manually await all tasks after invoking the target function by using the `wait_all` method. @@ -182,6 +186,18 @@ async def test_add_one(): # have been completed and do all the assertions. ``` +`wait_all()` also surfaces the first pending background callback failure even +if that callback completed before `wait_all()` was called. The failure is +consumed after it is raised, so a later `wait_all()` only observes newer work. +Cancelling a `wait_all()` call cancels only that waiter; accepted executions +remain tracked and can be drained by a later call. +Calling `shutdown()` performs the same drain before closing middleware, result +backend, and executor resources. If shutdown is cancelled while draining, it +finishes the drain and resource cleanup before propagating the cancellation. +Both drain methods must be called by the external test or application lifecycle +owner. Calling `wait_all()` or `shutdown()` from a task or `post_send` hook +managed by the same broker raises `RuntimeError` instead of waiting on itself. + ## Dependency injection If you use dependencies in your tasks, you may think that this can become a problem for tests. But it's not. diff --git a/taskiq/abc/broker.py b/taskiq/abc/broker.py index ea2e86c0..0eeb1fb8 100644 --- a/taskiq/abc/broker.py +++ b/taskiq/abc/broker.py @@ -23,7 +23,7 @@ from taskiq.acks import AckableMessage from taskiq.decor import AsyncTaskiqDecoratedTask from taskiq.events import TaskiqEvents -from taskiq.exceptions import TaskBrokerMismatchError +from taskiq.exceptions import SendTaskError, TaskBrokerMismatchError from taskiq.formatters.proxy_formatter import ProxyFormatter from taskiq.message import BrokerMessage from taskiq.result_backends.dummy import DummyResultBackend @@ -186,12 +186,9 @@ def add_middlewares(self, *middlewares: "TaskiqMiddleware") -> None: async def startup(self) -> None: """Do something when starting broker.""" - event = TaskiqEvents.CLIENT_STARTUP - if self.is_worker_process: - event = TaskiqEvents.WORKER_STARTUP - - for handler in self.event_handlers[event]: - await maybe_awaitable(handler(self.state)) + for event in self._get_startup_events(): + for handler in self.event_handlers[event]: + await maybe_awaitable(handler(self.state)) for middleware in self.middlewares: if middleware.__class__.startup != TaskiqMiddleware.startup: @@ -199,6 +196,12 @@ async def startup(self) -> None: await self.result_backend.startup() + def _get_startup_events(self) -> tuple[TaskiqEvents, ...]: + """Return event phases owned by this broker startup.""" + if self.is_worker_process: + return (TaskiqEvents.WORKER_STARTUP,) + return (TaskiqEvents.CLIENT_STARTUP,) + async def shutdown(self) -> None: """ Close the broker. @@ -206,19 +209,67 @@ async def shutdown(self) -> None: This method is called, when broker is closing. """ - event = TaskiqEvents.CLIENT_SHUTDOWN - if self.is_worker_process: - event = TaskiqEvents.WORKER_SHUTDOWN + shutdown_error: BaseException | None = None - # Call all shutdown events. - for handler in self.event_handlers[event]: - await maybe_awaitable(handler(self.state)) + for event in self._get_shutdown_events(): + for handler in self.event_handlers[event]: + try: + await maybe_awaitable(handler(self.state)) + except BaseException as exc: + shutdown_error = self._remember_shutdown_error( + shutdown_error, + exc, + ) for middleware in self.middlewares: if middleware.__class__.shutdown != TaskiqMiddleware.shutdown: - await maybe_awaitable(middleware.shutdown()) + try: + await maybe_awaitable(middleware.shutdown()) + except BaseException as exc: + shutdown_error = self._remember_shutdown_error( + shutdown_error, + exc, + ) + + try: + await self.result_backend.shutdown() + except BaseException as exc: + shutdown_error = self._remember_shutdown_error(shutdown_error, exc) + + if shutdown_error is not None: + raise shutdown_error + + def _get_shutdown_events(self) -> tuple[TaskiqEvents, ...]: + """Return event phases owned by this broker shutdown.""" + if self.is_worker_process: + return (TaskiqEvents.WORKER_SHUTDOWN,) + return (TaskiqEvents.CLIENT_SHUTDOWN,) + + @staticmethod + def _remember_shutdown_error( + first_error: BaseException | None, + current_error: BaseException, + ) -> BaseException: + """Keep the first shutdown failure while cleanup continues.""" + if first_error is None: + return current_error + logger.error( + "Additional error while shutting down broker resources.", + exc_info=current_error, + ) + return first_error - await self.result_backend.shutdown() + async def _kick_with_post_send( + self, + message: BrokerMessage, + post_send: Callable[[], Awaitable[None]], + ) -> None: + """Run the package-internal send boundary used by AsyncKicker.""" + try: + await self.kick(message) + except Exception as exc: + raise SendTaskError from exc + await post_send() @abstractmethod async def kick( diff --git a/taskiq/brokers/inmemory_broker.py b/taskiq/brokers/inmemory_broker.py index 0a7cc98e..5d82969a 100644 --- a/taskiq/brokers/inmemory_broker.py +++ b/taskiq/brokers/inmemory_broker.py @@ -1,6 +1,6 @@ import asyncio from collections import OrderedDict -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, Awaitable, Callable from concurrent.futures import ThreadPoolExecutor from typing import Any, TypeVar @@ -8,10 +8,9 @@ from taskiq.abc.result_backend import AsyncResultBackend, TaskiqResult from taskiq.depends.progress_tracker import TaskProgress from taskiq.events import TaskiqEvents -from taskiq.exceptions import UnknownTaskError +from taskiq.exceptions import SendTaskError, UnknownTaskError from taskiq.message import BrokerMessage from taskiq.receiver import Receiver -from taskiq.utils import maybe_awaitable _ReturnType = TypeVar("_ReturnType") @@ -145,7 +144,12 @@ def __init__( propagate_exceptions=propagate_exceptions, ) self.await_inplace = await_inplace + self._inflight_tasks: dict[ + asyncio.Task[Any], + asyncio.Task[Any] | None, + ] = {} self._running_tasks: set[asyncio.Task[Any]] = set() + self._running_task_error: BaseException | None = None async def kick(self, message: BrokerMessage) -> None: """ @@ -157,18 +161,136 @@ async def kick(self, message: BrokerMessage) -> None: :raises TaskiqError: if someone wants to kick unknown task. """ - target_task = self.find_task(message.task_name) - if target_task is None: + self._ensure_known_task(message) + await self._dispatch_message(message) + + async def _kick_with_post_send( + self, + message: BrokerMessage, + post_send: Callable[[], Awaitable[None]], + ) -> None: + """Run client post-send hooks before local task execution starts.""" + try: + self._ensure_known_task(message) + except Exception as exc: + raise SendTaskError from exc + + execution_task, execution_ready = self._create_gated_execution(message) + + try: + await post_send() + except asyncio.CancelledError: + self._release_to_background(execution_task, execution_ready) + raise + except Exception as post_send_error: + if not self.await_inplace: + self._release_to_background(execution_task, execution_ready) + raise + await self._finish_after_post_send_error( + execution_task, + execution_ready, + post_send_error, + ) + except BaseException: + self._release_to_background(execution_task, execution_ready) + raise + + if not self.await_inplace: + self._release_to_background(execution_task, execution_ready) + return + + execution_ready.set() + try: + await execution_task + except Exception as exc: + raise SendTaskError from exc + + def _create_gated_execution( + self, + message: BrokerMessage, + ) -> tuple[asyncio.Task[Any], asyncio.Event]: + """Create accepted execution that waits for the send boundary.""" + execution_ready = asyncio.Event() + execution_task = asyncio.create_task( + self._execute_after_post_send(message, execution_ready), + ) + self._track_inflight_task(execution_task) + return execution_task, execution_ready + + async def _execute_after_post_send( + self, + message: BrokerMessage, + execution_ready: asyncio.Event, + ) -> None: + """Wait for post-send completion before executing a local message.""" + await execution_ready.wait() + await self.receiver.callback(message=message.message) + + async def _finish_after_post_send_error( + self, + task: asyncio.Task[Any], + execution_ready: asyncio.Event, + post_send_error: Exception, + ) -> None: + """Finish accepted inline work without replacing its post-send error.""" + execution_ready.set() + try: + await task + except asyncio.CancelledError: + raise + except BaseException as execution_error: + raise post_send_error from execution_error + raise post_send_error + + def _ensure_known_task(self, message: BrokerMessage) -> None: + """Reject direct sends for tasks that are not registered locally.""" + if self.find_task(message.task_name) is None: raise UnknownTaskError(task_name=message.task_name) - receiver_cb = self.receiver.callback(message=message.message) + async def _dispatch_message(self, message: BrokerMessage) -> None: + """Execute a validated message inline or track it in the background.""" if self.await_inplace: - await receiver_cb + task = asyncio.create_task( + self.receiver.callback(message=message.message), + ) + self._track_inflight_task(task) + await task return - task = asyncio.create_task(receiver_cb) + self._start_background_task(message) + + def _start_background_task(self, message: BrokerMessage) -> None: + """Start and track one local receiver callback.""" + task = asyncio.create_task(self.receiver.callback(message=message.message)) + self._track_running_task(task) + + def _track_running_task(self, task: asyncio.Task[Any]) -> None: + """Track one local receiver callback until it finishes.""" self._running_tasks.add(task) - task.add_done_callback(self._running_tasks.discard) + task.add_done_callback(self._on_running_task_done) + + def _track_inflight_task(self, task: asyncio.Task[Any]) -> None: + """Track execution owned by an active inline send.""" + self._inflight_tasks[task] = asyncio.current_task() + task.add_done_callback(self._on_inflight_task_done) + + def _on_inflight_task_done(self, task: asyncio.Task[Any]) -> None: + """Release execution after its sender has observed the result.""" + self._inflight_tasks.pop(task, None) + + def _promote_to_background(self, task: asyncio.Task[Any]) -> None: + """Transfer execution ownership from the sender to broker drain.""" + self._inflight_tasks.pop(task, None) + self._track_running_task(task) + + def _release_to_background( + self, + task: asyncio.Task[Any], + execution_ready: asyncio.Event, + ) -> None: + """Transfer and release gated execution as one ownership transition.""" + self._promote_to_background(task) + execution_ready.set() def listen(self) -> AsyncGenerator[bytes, None]: """ @@ -187,20 +309,95 @@ async def wait_all(self) -> None: Useful when used in testing and you need to await all sent tasks before asserting results. + Cancelling this waiter does not cancel accepted executions. """ - to_await = list(self._running_tasks) - for task in to_await: - await task + self._ensure_external_drain("wait_all") + + while self._inflight_tasks or self._running_tasks: + active_tasks = tuple(self._inflight_tasks.keys() | self._running_tasks) + await asyncio.gather( + *(asyncio.shield(task) for task in active_tasks), + return_exceptions=True, + ) + for task in active_tasks: + self._inflight_tasks.pop(task, None) + self._on_running_task_done(task) + + if self._running_task_error is not None: + task_error = self._running_task_error + self._running_task_error = None + raise task_error + + def _on_running_task_done(self, task: asyncio.Task[Any]) -> None: + """Release a completed task and retain its first failure until drain.""" + if task not in self._running_tasks: + return + + self._running_tasks.remove(task) + try: + task.result() + except BaseException as exc: + if self._running_task_error is None: + self._running_task_error = exc + + def _ensure_external_drain(self, operation: str) -> None: + """Reject drain calls made by work that the drain must await.""" + current_task = asyncio.current_task() + if current_task is None: + return + + is_managed_execution = ( + current_task in self._running_tasks + or current_task in self._inflight_tasks + or current_task in self._inflight_tasks.values() + ) + if is_managed_execution: + raise RuntimeError( + f"InMemoryBroker.{operation}() cannot be called from " + "a task managed by this broker.", + ) + + def _get_startup_events(self) -> tuple[TaskiqEvents, ...]: + """Run both sides because tasks execute in the client process.""" + return (TaskiqEvents.CLIENT_STARTUP, TaskiqEvents.WORKER_STARTUP) - async def startup(self) -> None: - """Runs startup events for client and worker side.""" - for event in (TaskiqEvents.CLIENT_STARTUP, TaskiqEvents.WORKER_STARTUP): - for handler in self.event_handlers.get(event, []): - await maybe_awaitable(handler(self.state)) + def _get_shutdown_events(self) -> tuple[TaskiqEvents, ...]: + """Shut down both event phases in their legacy order.""" + return (TaskiqEvents.CLIENT_SHUTDOWN, TaskiqEvents.WORKER_SHUTDOWN) async def shutdown(self) -> None: - """Runs shutdown events for client and worker side.""" - for event in (TaskiqEvents.CLIENT_SHUTDOWN, TaskiqEvents.WORKER_SHUTDOWN): - for handler in self.event_handlers.get(event, []): - await maybe_awaitable(handler(self.state)) - self.executor.shutdown() + """Drain local execution, close resources and stop the executor. + + Cancellation is propagated after accepted work and resources are closed. + """ + self._ensure_external_drain("shutdown") + shutdown_error: BaseException | None = None + + while True: + try: + await self.wait_all() + except asyncio.CancelledError as exc: + shutdown_error = self._remember_shutdown_error( + shutdown_error, + exc, + ) + continue + except BaseException as exc: + shutdown_error = self._remember_shutdown_error( + shutdown_error, + exc, + ) + break + + try: + await super().shutdown() + except BaseException as exc: + shutdown_error = self._remember_shutdown_error(shutdown_error, exc) + + try: + self.executor.shutdown() + except BaseException as exc: # pragma: no cover + shutdown_error = self._remember_shutdown_error(shutdown_error, exc) + + if shutdown_error is not None: + raise shutdown_error diff --git a/taskiq/kicker.py b/taskiq/kicker.py index b8b9c073..e82945de 100644 --- a/taskiq/kicker.py +++ b/taskiq/kicker.py @@ -1,6 +1,7 @@ from collections.abc import Coroutine from dataclasses import asdict, is_dataclass from datetime import datetime, timedelta +from functools import partial from logging import getLogger from types import CoroutineType from typing import ( @@ -141,10 +142,10 @@ async def kiq( **kwargs: _FuncParams.kwargs, ) -> Any: """ - This method sends function call over the network. + Prepare and dispatch a task invocation through the current broker. - It gets current broker and calls it's kick method, - returning what it returns. + Send middleware runs around broker dispatch before the task handle is + returned. :param args: function's arguments. :param kwargs: function's key word arguments. @@ -160,14 +161,17 @@ async def kiq( for middleware in self.broker.middlewares: if middleware.__class__.pre_send != TaskiqMiddleware.pre_send: message = await maybe_awaitable(middleware.pre_send(message)) + try: - await self.broker.kick(self.broker.formatter.dumps(message)) + broker_message = self.broker.formatter.dumps(message) except Exception as exc: raise SendTaskError from exc - for middleware in reversed(self.broker.middlewares): - if middleware.__class__.post_send != TaskiqMiddleware.post_send: - await maybe_awaitable(middleware.post_send(message)) + # AsyncKicker and AsyncBroker share this package-internal send boundary. + await self.broker._kick_with_post_send( # noqa: SLF001 + broker_message, + partial(self._run_post_send, message), + ) return AsyncTaskiqTask( task_id=message.task_id, @@ -175,6 +179,12 @@ async def kiq( return_type=self.return_type, # type: ignore # (pyright issue) ) + async def _run_post_send(self, message: TaskiqMessage) -> None: + """Run post-send middleware in reverse registration order.""" + for middleware in reversed(self.broker.middlewares): + if middleware.__class__.post_send != TaskiqMiddleware.post_send: + await maybe_awaitable(middleware.post_send(message)) + async def schedule_by_cron( self, source: "ScheduleSource", diff --git a/tests/abc/test_broker.py b/tests/abc/test_broker.py index 636f9576..8ce26da0 100644 --- a/tests/abc/test_broker.py +++ b/tests/abc/test_broker.py @@ -4,9 +4,11 @@ import pytest from taskiq.abc.broker import AsyncBroker +from taskiq.abc.middleware import TaskiqMiddleware from taskiq.decor import AsyncTaskiqDecoratedTask from taskiq.events import TaskiqEvents -from taskiq.message import BrokerMessage +from taskiq.exceptions import SendTaskError +from taskiq.message import BrokerMessage, TaskiqMessage from taskiq.state import TaskiqState @@ -30,6 +32,46 @@ async def listen(self) -> AsyncGenerator[BrokerMessage, None]: # type: ignore """ +class _RecordingSendBroker(_TestBroker): + """Record transport dispatch and raise an optional send failure.""" + + def __init__( + self, + events: list[str], + kick_error: Exception | None = None, + ) -> None: + super().__init__() + self.events = events + self.kick_error = kick_error + + async def kick(self, message: BrokerMessage) -> None: + self.events.append("kick") + if self.kick_error is not None: + raise self.kick_error + + +class _RecordingSendMiddleware(TaskiqMiddleware): + """Record client send hooks and raise an optional post-send failure.""" + + def __init__( + self, + events: list[str], + post_send_error: Exception | None = None, + ) -> None: + super().__init__() + self.events = events + self.post_send_error = post_send_error + + def pre_send(self, message: TaskiqMessage) -> TaskiqMessage: + self.events.append("pre_send") + return message + + def post_send(self, message: TaskiqMessage) -> None: + self.events.append("post_send") + if self.post_send_error is not None: + raise self.post_send_error + + def test_decorator_success() -> None: """Test that decoration without parameters works.""" tbrok = _TestBroker() @@ -82,6 +124,56 @@ async def test_task() -> None: ... assert test_task.labels == old_labels +async def test_kicker_preserves_external_broker_send_order() -> None: + events: list[str] = [] + broker = _RecordingSendBroker(events) + broker.with_middlewares(_RecordingSendMiddleware(events)) + + @broker.task + async def task() -> None: + return None + + await task.kiq() + + assert events == ["pre_send", "kick", "post_send"] + + +async def test_kicker_wraps_external_broker_transport_error() -> None: + events: list[str] = [] + kick_error = RuntimeError("transport failed") + broker = _RecordingSendBroker(events, kick_error=kick_error) + broker.with_middlewares(_RecordingSendMiddleware(events)) + + @broker.task + async def task() -> None: + return None + + with pytest.raises(SendTaskError) as exc_info: + await task.kiq() + + assert exc_info.value.__cause__ is kick_error + assert events == ["pre_send", "kick"] + + +async def test_kicker_preserves_post_send_error_type() -> None: + events: list[str] = [] + post_send_error = RuntimeError("post-send failed") + broker = _RecordingSendBroker(events) + broker.with_middlewares( + _RecordingSendMiddleware(events, post_send_error=post_send_error), + ) + + @broker.task + async def task() -> None: + return None + + with pytest.raises(RuntimeError) as exc_info: + await task.kiq() + + assert exc_info.value is post_send_error + assert events == ["pre_send", "kick", "post_send"] + + @pytest.mark.anyio @pytest.mark.parametrize( ("is_worker_process", "startup", "shutdown"), diff --git a/tests/brokers/inmemory_contract_support.py b/tests/brokers/inmemory_contract_support.py new file mode 100644 index 00000000..1fe6af33 --- /dev/null +++ b/tests/brokers/inmemory_contract_support.py @@ -0,0 +1,182 @@ +import asyncio +from typing import Any + +from taskiq import InMemoryBroker, TaskiqMessage, TaskiqResult +from taskiq.abc.middleware import TaskiqMiddleware +from taskiq.brokers.inmemory_broker import InmemoryResultBackend + + +class LifecycleError(Exception): + """Marker exception for lifecycle fault tests.""" + + +class PostSendError(Exception): + """Marker exception for post-send compatibility tests.""" + + +class RecordingLifecycleMiddleware(TaskiqMiddleware): + """Record middleware lifecycle calls and optional shutdown failure.""" + + def __init__( + self, + events: list[str], + shutdown_error: Exception | None = None, + ) -> None: + super().__init__() + self.events = events + self.shutdown_error = shutdown_error + + async def startup(self) -> None: + self.events.append("middleware.startup") + + async def shutdown(self) -> None: + self.events.append("middleware.shutdown") + if self.shutdown_error is not None: + raise self.shutdown_error + + +class RecordingResultBackend(InmemoryResultBackend[Any]): + """Record result backend lifecycle calls and optional failures.""" + + def __init__( + self, + events: list[str], + startup_error: Exception | None = None, + shutdown_error: Exception | None = None, + ) -> None: + super().__init__() + self.events = events + self.startup_error = startup_error + self.shutdown_error = shutdown_error + + async def startup(self) -> None: + self.events.append("backend.startup") + if self.startup_error is not None: + raise self.startup_error + + async def shutdown(self) -> None: + self.events.append("backend.shutdown") + if self.shutdown_error is not None: + raise self.shutdown_error + + +class FailingExecutionMiddleware(TaskiqMiddleware): + """Raise before execution to expose background callback failures.""" + + def __init__(self, error: Exception) -> None: + super().__init__() + self.error = error + + def pre_execute(self, message: TaskiqMessage) -> TaskiqMessage: + raise self.error + + +class BlockingPostSendMiddleware(TaskiqMiddleware): + """Pause post-send to prove execution cannot overlap it.""" + + def __init__( + self, + events: list[str], + post_send_started: asyncio.Event, + release_post_send: asyncio.Event, + ) -> None: + super().__init__() + self.events = events + self.post_send_started = post_send_started + self.release_post_send = release_post_send + + def pre_send(self, message: TaskiqMessage) -> TaskiqMessage: + self.events.append("pre_send") + return message + + async def post_send(self, message: TaskiqMessage) -> None: + self.events.append("post_send.started") + self.post_send_started.set() + await self.release_post_send.wait() + self.events.append("post_send.finished") + + def pre_execute(self, message: TaskiqMessage) -> TaskiqMessage: + self.events.append("pre_execute") + return message + + def post_execute( + self, + message: TaskiqMessage, + result: TaskiqResult[Any], + ) -> None: + self.events.append("post_execute") + + +class CoordinatedPostSendMiddleware(TaskiqMiddleware): + """Hold two post-send hooks at different deterministic boundaries.""" + + def __init__( + self, + first_task_name: str, + second_task_name: str, + first_post_send_started: asyncio.Event, + second_post_send_started: asyncio.Event, + release_second_post_send: asyncio.Event, + ) -> None: + super().__init__() + self.first_task_name = first_task_name + self.second_task_name = second_task_name + self.first_post_send_started = first_post_send_started + self.second_post_send_started = second_post_send_started + self.release_second_post_send = release_second_post_send + + async def post_send(self, message: TaskiqMessage) -> None: + if message.task_name == self.first_task_name: + self.first_post_send_started.set() + await self.second_post_send_started.wait() + return + + if message.task_name == self.second_task_name: + self.second_post_send_started.set() + await self.release_second_post_send.wait() + + +class FailingPostSendMiddleware(TaskiqMiddleware): + """Fail after transport acceptance.""" + + def __init__(self, error: Exception) -> None: + super().__init__() + self.error = error + + def post_send(self, message: TaskiqMessage) -> None: + raise self.error + + +class ReentrantDrainMiddleware(TaskiqMiddleware): + """Attempt a broker drain from the active send lifecycle.""" + + def __init__(self, operation: str) -> None: + super().__init__() + self.operation = operation + + async def post_send(self, message: TaskiqMessage) -> None: + await getattr(self.broker, self.operation)() + + +class DrainSignallingInMemoryBroker(InMemoryBroker): + """Expose drain attempts for deterministic overlap tests.""" + + def __init__( + self, + drain_started: asyncio.Event, + drain_restarted: asyncio.Event | None = None, + *, + await_inplace: bool = False, + ) -> None: + super().__init__(await_inplace=await_inplace) + self.drain_started = drain_started + self.drain_restarted = drain_restarted + self.drain_calls = 0 + + async def wait_all(self) -> None: + self.drain_calls += 1 + if self.drain_calls == 1: + self.drain_started.set() + elif self.drain_restarted is not None: + self.drain_restarted.set() + await super().wait_all() diff --git a/tests/brokers/test_inmemory_lifecycle.py b/tests/brokers/test_inmemory_lifecycle.py new file mode 100644 index 00000000..2a33d388 --- /dev/null +++ b/tests/brokers/test_inmemory_lifecycle.py @@ -0,0 +1,301 @@ +import asyncio +import logging +from typing import Any + +import pytest + +from taskiq import InMemoryBroker +from taskiq.events import TaskiqEvents +from taskiq.state import TaskiqState +from tests.brokers.inmemory_contract_support import ( + DrainSignallingInMemoryBroker, + FailingExecutionMiddleware, + LifecycleError, + RecordingLifecycleMiddleware, + RecordingResultBackend, +) + + +async def test_lifecycle_runs_events_and_resources_once_in_order() -> None: + events: list[str] = [] + broker = InMemoryBroker() + broker.with_middlewares(RecordingLifecycleMiddleware(events)) + broker.with_result_backend(RecordingResultBackend(events)) + + @broker.on_event(TaskiqEvents.CLIENT_STARTUP) + def record_client_startup(state: TaskiqState) -> None: + events.append("client.startup") + + @broker.on_event(TaskiqEvents.WORKER_STARTUP) + def record_worker_startup(state: TaskiqState) -> None: + events.append("worker.startup") + + @broker.on_event(TaskiqEvents.CLIENT_SHUTDOWN) + def record_client_shutdown(state: TaskiqState) -> None: + events.append("client.shutdown") + + @broker.on_event(TaskiqEvents.WORKER_SHUTDOWN) + def record_worker_shutdown(state: TaskiqState) -> None: + events.append("worker.shutdown") + + await broker.startup() + await broker.shutdown() + + assert events == [ + "client.startup", + "worker.startup", + "middleware.startup", + "backend.startup", + "client.shutdown", + "worker.shutdown", + "middleware.shutdown", + "backend.shutdown", + ] + + +async def test_shutdown_cleans_resources_after_startup_failure() -> None: + events: list[str] = [] + startup_error = LifecycleError("backend startup failed") + broker = InMemoryBroker() + broker.with_middlewares(RecordingLifecycleMiddleware(events)) + broker.with_result_backend(RecordingResultBackend(events, startup_error)) + + with pytest.raises(LifecycleError) as exc_info: + await broker.startup() + assert exc_info.value is startup_error + + await broker.shutdown() + + assert events == [ + "middleware.startup", + "backend.startup", + "middleware.shutdown", + "backend.shutdown", + ] + + +async def test_shutdown_keeps_first_failure_and_closes_every_resource( + caplog: pytest.LogCaptureFixture, +) -> None: + events: list[str] = [] + event_error = LifecycleError("event shutdown failed") + middleware_error = LifecycleError("middleware shutdown failed") + backend_error = LifecycleError("backend shutdown failed") + broker = InMemoryBroker() + broker.with_middlewares( + RecordingLifecycleMiddleware(events, shutdown_error=middleware_error), + ) + broker.with_result_backend( + RecordingResultBackend(events, shutdown_error=backend_error), + ) + + @broker.on_event(TaskiqEvents.CLIENT_SHUTDOWN) + def fail_client_shutdown(state: TaskiqState) -> None: + events.append("client.shutdown") + raise event_error + + @broker.on_event(TaskiqEvents.WORKER_SHUTDOWN) + def record_worker_shutdown(state: TaskiqState) -> None: + events.append("worker.shutdown") + + caplog.set_level(logging.ERROR, logger="taskiq") + await broker.startup() + + with pytest.raises(LifecycleError) as exc_info: + await broker.shutdown() + + assert exc_info.value is event_error + assert events == [ + "middleware.startup", + "backend.startup", + "client.shutdown", + "worker.shutdown", + "middleware.shutdown", + "backend.shutdown", + ] + assert caplog.text.count("Additional error while shutting down") == 2 + with pytest.raises(RuntimeError, match="cannot schedule new futures"): + broker.executor.submit(int) + + +async def test_shutdown_cancellation_still_closes_every_resource() -> None: + events: list[str] = [] + shutdown_started = asyncio.Event() + keep_shutdown_blocked = asyncio.Event() + broker = InMemoryBroker() + broker.with_middlewares(RecordingLifecycleMiddleware(events)) + broker.with_result_backend(RecordingResultBackend(events)) + + @broker.on_event(TaskiqEvents.CLIENT_SHUTDOWN) + async def block_client_shutdown(state: TaskiqState) -> None: + events.append("client.shutdown") + shutdown_started.set() + await keep_shutdown_blocked.wait() + + await broker.startup() + shutdown_task = asyncio.create_task(broker.shutdown()) + await asyncio.wait_for(shutdown_started.wait(), timeout=1) + shutdown_task.cancel() + + with pytest.raises(asyncio.CancelledError): + await shutdown_task + + assert events == [ + "middleware.startup", + "backend.startup", + "client.shutdown", + "middleware.shutdown", + "backend.shutdown", + ] + with pytest.raises(RuntimeError, match="cannot schedule new futures"): + broker.executor.submit(int) + + +async def test_cancelled_shutdown_drains_before_resource_cleanup() -> None: + task_started = asyncio.Event() + finish_task = asyncio.Event() + drain_started = asyncio.Event() + drain_restarted = asyncio.Event() + events: list[str] = [] + broker = DrainSignallingInMemoryBroker(drain_started, drain_restarted) + broker.with_middlewares(RecordingLifecycleMiddleware(events)) + broker.with_result_backend(RecordingResultBackend(events)) + + @broker.task + async def running_task() -> None: + events.append("task.started") + task_started.set() + await finish_task.wait() + events.append("task.finished") + + await broker.startup() + await running_task.kiq() + await asyncio.wait_for(task_started.wait(), timeout=1) + + shutdown_task = asyncio.create_task(broker.shutdown()) + await asyncio.wait_for(drain_started.wait(), timeout=1) + shutdown_task.cancel() + await asyncio.wait_for(drain_restarted.wait(), timeout=1) + + assert not shutdown_task.done() + finish_task.set() + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(asyncio.shield(shutdown_task), timeout=1) + + assert events == [ + "middleware.startup", + "backend.startup", + "task.started", + "task.finished", + "middleware.shutdown", + "backend.shutdown", + ] + assert not broker._running_tasks + with pytest.raises(RuntimeError, match="cannot schedule new futures"): + broker.executor.submit(int) + + +async def test_shutdown_waits_for_running_tasks_before_resource_cleanup() -> None: + task_started = asyncio.Event() + finish_task = asyncio.Event() + drain_started = asyncio.Event() + events: list[str] = [] + broker = DrainSignallingInMemoryBroker(drain_started) + broker.with_middlewares(RecordingLifecycleMiddleware(events)) + broker.with_result_backend(RecordingResultBackend(events)) + + @broker.on_event(TaskiqEvents.CLIENT_SHUTDOWN) + def record_client_shutdown(state: TaskiqState) -> None: + events.append("client.shutdown") + + @broker.task + async def running_task() -> None: + events.append("task.started") + task_started.set() + await finish_task.wait() + events.append("task.finished") + + await broker.startup() + await running_task.kiq() + await asyncio.wait_for(task_started.wait(), timeout=1) + + shutdown_task = asyncio.create_task(broker.shutdown()) + await asyncio.wait_for(drain_started.wait(), timeout=1) + assert not shutdown_task.done() + + finish_task.set() + await asyncio.wait_for(shutdown_task, timeout=1) + + assert events == [ + "middleware.startup", + "backend.startup", + "task.started", + "task.finished", + "client.shutdown", + "middleware.shutdown", + "backend.shutdown", + ] + assert not broker._running_tasks + + +async def test_wait_all_retains_failure_from_already_completed_task() -> None: + task_error = LifecycleError("completed in-memory task failed") + broker = InMemoryBroker() + broker.with_middlewares(FailingExecutionMiddleware(task_error)) + + @broker.task + async def failing_task() -> None: + return None + + await failing_task.kiq() + running_task = next(iter(broker._running_tasks)) + task_completed = asyncio.Event() + + def mark_completed(completed_task: asyncio.Task[Any]) -> None: + assert completed_task.done() + task_completed.set() + + running_task.add_done_callback(mark_completed) + await asyncio.wait_for(task_completed.wait(), timeout=1) + assert not broker._running_tasks + + with pytest.raises(LifecycleError) as exc_info: + await broker.wait_all() + + assert exc_info.value is task_error + await broker.wait_all() + await broker.shutdown() + + +async def test_shutdown_cleans_resources_after_task_failure() -> None: + events: list[str] = [] + task_error = LifecycleError("in-memory task failed") + middleware_error = LifecycleError("middleware shutdown failed") + broker = InMemoryBroker() + broker.with_middlewares(FailingExecutionMiddleware(task_error)) + broker.with_middlewares( + RecordingLifecycleMiddleware(events, shutdown_error=middleware_error), + ) + broker.with_result_backend(RecordingResultBackend(events)) + + @broker.task + async def failing_task() -> None: + return None + + await broker.startup() + await failing_task.kiq() + + with pytest.raises(LifecycleError) as exc_info: + await broker.shutdown() + + assert exc_info.value is task_error + assert events == [ + "middleware.startup", + "backend.startup", + "middleware.shutdown", + "backend.shutdown", + ] + assert not broker._running_tasks + with pytest.raises(RuntimeError, match="cannot schedule new futures"): + broker.executor.submit(int) diff --git a/tests/brokers/test_inmemory_send.py b/tests/brokers/test_inmemory_send.py new file mode 100644 index 00000000..690ca7aa --- /dev/null +++ b/tests/brokers/test_inmemory_send.py @@ -0,0 +1,400 @@ +import asyncio +from typing import Any + +import pytest + +from taskiq import InMemoryBroker, TaskiqMessage +from taskiq.exceptions import SendTaskError, UnknownTaskError +from taskiq.kicker import AsyncKicker +from tests.brokers.inmemory_contract_support import ( + BlockingPostSendMiddleware, + CoordinatedPostSendMiddleware, + DrainSignallingInMemoryBroker, + FailingExecutionMiddleware, + FailingPostSendMiddleware, + LifecycleError, + PostSendError, + ReentrantDrainMiddleware, +) + + +async def test_direct_kick_preserves_inline_execution() -> None: + task_executed = False + broker = InMemoryBroker(await_inplace=True) + + @broker.task(task_name="direct.task") + async def task() -> None: + nonlocal task_executed + task_executed = True + + message = TaskiqMessage( + task_id="direct-task-id", + task_name=task.task_name, + labels={}, + labels_types={}, + args=[], + kwargs={}, + ) + await broker.kick(broker.formatter.dumps(message)) + + assert task_executed + await broker.wait_all() + await broker.shutdown() + + +@pytest.mark.parametrize("await_inplace", [False, True]) +async def test_async_post_send_finishes_before_execution( + await_inplace: bool, +) -> None: + events: list[str] = [] + post_send_started = asyncio.Event() + release_post_send = asyncio.Event() + task_started = asyncio.Event() + broker = InMemoryBroker(await_inplace=await_inplace) + broker.with_middlewares( + BlockingPostSendMiddleware( + events, + post_send_started, + release_post_send, + ), + ) + + @broker.task + async def task() -> None: + events.append("task") + task_started.set() + + kiq_task = asyncio.create_task(task.kiq()) + await asyncio.wait_for(post_send_started.wait(), timeout=1) + + assert not task_started.is_set() + assert events == ["pre_send", "post_send.started"] + + release_post_send.set() + await asyncio.wait_for(kiq_task, timeout=1) + await asyncio.wait_for(task_started.wait(), timeout=1) + await broker.shutdown() + + assert events == [ + "pre_send", + "post_send.started", + "post_send.finished", + "pre_execute", + "task", + "post_execute", + ] + + +async def test_concurrent_inline_sends_keep_per_invocation_ownership() -> None: + first_post_send_started = asyncio.Event() + second_post_send_started = asyncio.Event() + release_second_post_send = asyncio.Event() + first_task_finished = asyncio.Event() + second_task_started = asyncio.Event() + finish_second_task = asyncio.Event() + second_task_finished = asyncio.Event() + broker = InMemoryBroker(await_inplace=True) + + @broker.task(task_name="first.task") + async def first_task() -> None: + first_task_finished.set() + + @broker.task(task_name="second.task") + async def second_task() -> None: + second_task_started.set() + await finish_second_task.wait() + second_task_finished.set() + + broker.with_middlewares( + CoordinatedPostSendMiddleware( + first_task.task_name, + second_task.task_name, + first_post_send_started, + second_post_send_started, + release_second_post_send, + ), + ) + + first_sender = asyncio.create_task(first_task.kiq()) + await asyncio.wait_for(first_post_send_started.wait(), timeout=1) + second_sender = asyncio.create_task(second_task.kiq()) + await asyncio.wait_for(second_post_send_started.wait(), timeout=1) + + await asyncio.wait_for(first_sender, timeout=1) + assert first_task_finished.is_set() + assert not second_task_started.is_set() + assert not second_sender.done() + + release_second_post_send.set() + await asyncio.wait_for(second_task_started.wait(), timeout=1) + assert not second_sender.done() + + finish_second_task.set() + await asyncio.wait_for(second_sender, timeout=1) + + assert second_task_finished.is_set() + await broker.shutdown() + + +async def test_post_send_failure_does_not_retract_accepted_inline_task() -> None: + post_send_error = PostSendError("post-send failed") + task_executed = False + broker = InMemoryBroker(await_inplace=True) + broker.with_middlewares(FailingPostSendMiddleware(post_send_error)) + + @broker.task + async def task() -> None: + nonlocal task_executed + task_executed = True + + with pytest.raises(PostSendError) as exc_info: + await task.kiq() + + assert exc_info.value is post_send_error + assert task_executed + await broker.shutdown() + + +async def test_inline_execution_failure_remains_a_send_error() -> None: + execution_error = LifecycleError("execution failed") + broker = InMemoryBroker(await_inplace=True) + broker.with_middlewares(FailingExecutionMiddleware(execution_error)) + + @broker.task + async def task() -> None: + return None + + with pytest.raises(SendTaskError) as exc_info: + await task.kiq() + + assert exc_info.value.__cause__ is execution_error + await broker.wait_all() + await broker.shutdown() + + +async def test_kicker_wraps_unknown_inmemory_task() -> None: + broker = InMemoryBroker() + kicker: AsyncKicker[Any, Any] = AsyncKicker("missing.task", broker, {}) + + with pytest.raises(SendTaskError) as exc_info: + await kicker.kiq() + + assert isinstance(exc_info.value.__cause__, UnknownTaskError) + await broker.shutdown() + + +@pytest.mark.parametrize("await_inplace", [False, True]) +async def test_post_send_failure_preserves_execution_error_ownership( + await_inplace: bool, +) -> None: + post_send_error = PostSendError("post-send failed") + execution_error = LifecycleError("execution failed") + broker = InMemoryBroker(await_inplace=await_inplace) + broker.with_middlewares( + FailingPostSendMiddleware(post_send_error), + FailingExecutionMiddleware(execution_error), + ) + + @broker.task + async def task() -> None: + return None + + with pytest.raises(PostSendError) as exc_info: + await task.kiq() + + assert exc_info.value is post_send_error + if await_inplace: + assert exc_info.value.__cause__ is execution_error + await broker.wait_all() + else: + with pytest.raises(LifecycleError) as execution_exc_info: + await broker.wait_all() + assert execution_exc_info.value is execution_error + + await broker.shutdown() + + +async def test_post_send_cancellation_does_not_retract_accepted_inline_task() -> None: + events: list[str] = [] + post_send_started = asyncio.Event() + release_post_send = asyncio.Event() + task_executed = asyncio.Event() + broker = InMemoryBroker(await_inplace=True) + broker.with_middlewares( + BlockingPostSendMiddleware( + events, + post_send_started, + release_post_send, + ), + ) + + @broker.task + async def task() -> None: + task_executed.set() + + kiq_task = asyncio.create_task(task.kiq()) + await asyncio.wait_for(post_send_started.wait(), timeout=1) + kiq_task.cancel() + + with pytest.raises(asyncio.CancelledError): + await kiq_task + + await asyncio.wait_for(task_executed.wait(), timeout=1) + await asyncio.wait_for(broker.wait_all(), timeout=1) + await broker.shutdown() + + +async def test_sender_cancellation_preserves_inline_execution_cancellation() -> None: + task_started = asyncio.Event() + task_cancelled = asyncio.Event() + keep_running = asyncio.Event() + broker = InMemoryBroker(await_inplace=True) + + @broker.task + async def task() -> None: + task_started.set() + try: + await keep_running.wait() + except asyncio.CancelledError: + task_cancelled.set() + raise + + kiq_task = asyncio.create_task(task.kiq()) + await asyncio.wait_for(task_started.wait(), timeout=1) + kiq_task.cancel() + + await asyncio.wait_for(asyncio.shield(kiq_task), timeout=1) + + assert not kiq_task.cancelled() + assert task_cancelled.is_set() + await broker.wait_all() + await broker.shutdown() + + +async def test_shutdown_waits_for_send_blocked_in_post_send() -> None: + events: list[str] = [] + post_send_started = asyncio.Event() + release_post_send = asyncio.Event() + task_executed = asyncio.Event() + drain_started = asyncio.Event() + broker = DrainSignallingInMemoryBroker(drain_started) + broker.with_middlewares( + BlockingPostSendMiddleware( + events, + post_send_started, + release_post_send, + ), + ) + + @broker.task + async def task() -> None: + task_executed.set() + + kiq_task = asyncio.create_task(task.kiq()) + await asyncio.wait_for(post_send_started.wait(), timeout=1) + shutdown_task = asyncio.create_task(broker.shutdown()) + await asyncio.wait_for(drain_started.wait(), timeout=1) + + assert not shutdown_task.done() + assert not task_executed.is_set() + + release_post_send.set() + await kiq_task + await asyncio.wait_for(shutdown_task, timeout=1) + + assert task_executed.is_set() + + +@pytest.mark.parametrize("await_inplace", [False, True]) +async def test_cancelled_wait_all_does_not_cancel_accepted_send( + await_inplace: bool, +) -> None: + events: list[str] = [] + post_send_started = asyncio.Event() + release_post_send = asyncio.Event() + task_executed = asyncio.Event() + drain_started = asyncio.Event() + broker = DrainSignallingInMemoryBroker( + drain_started, + await_inplace=await_inplace, + ) + broker.with_middlewares( + BlockingPostSendMiddleware( + events, + post_send_started, + release_post_send, + ), + ) + + @broker.task + async def task() -> None: + task_executed.set() + + sender = asyncio.create_task(task.kiq()) + await asyncio.wait_for(post_send_started.wait(), timeout=1) + drain = asyncio.create_task(broker.wait_all()) + await asyncio.wait_for(drain_started.wait(), timeout=1) + drain.cancel() + + with pytest.raises(asyncio.CancelledError): + await drain + + assert not sender.done() + assert not task_executed.is_set() + + release_post_send.set() + await asyncio.wait_for(sender, timeout=1) + await asyncio.wait_for(task_executed.wait(), timeout=1) + await asyncio.wait_for(broker.wait_all(), timeout=1) + await broker.shutdown() + + +@pytest.mark.parametrize("operation", ["wait_all", "shutdown"]) +@pytest.mark.parametrize("await_inplace", [False, True]) +async def test_execution_cannot_reenter_broker_drain( + operation: str, + await_inplace: bool, +) -> None: + broker = InMemoryBroker(await_inplace=await_inplace) + drain_error: RuntimeError | None = None + + @broker.task + async def task() -> None: + nonlocal drain_error + try: + await getattr(broker, operation)() + except RuntimeError as exc: + drain_error = exc + + await asyncio.wait_for(task.kiq(), timeout=1) + await asyncio.wait_for(broker.wait_all(), timeout=1) + + assert drain_error is not None + assert str(drain_error).startswith(f"InMemoryBroker.{operation}()") + assert broker.executor.submit(int).result() == 0 + await broker.shutdown() + + +@pytest.mark.parametrize("operation", ["wait_all", "shutdown"]) +@pytest.mark.parametrize("await_inplace", [False, True]) +async def test_post_send_cannot_reenter_broker_drain( + operation: str, + await_inplace: bool, +) -> None: + task_executed = False + broker = InMemoryBroker(await_inplace=await_inplace) + broker.with_middlewares(ReentrantDrainMiddleware(operation)) + + @broker.task + async def task() -> None: + nonlocal task_executed + task_executed = True + + with pytest.raises(RuntimeError) as exc_info: + await asyncio.wait_for(task.kiq(), timeout=1) + + await asyncio.wait_for(broker.wait_all(), timeout=1) + assert str(exc_info.value).startswith(f"InMemoryBroker.{operation}()") + assert task_executed + assert broker.executor.submit(int).result() == 0 + await broker.shutdown() diff --git a/tests/middlewares/admin_middleware/conftest.py b/tests/middlewares/admin_middleware/conftest.py index dbec77a4..9375a9d2 100644 --- a/tests/middlewares/admin_middleware/conftest.py +++ b/tests/middlewares/admin_middleware/conftest.py @@ -39,13 +39,12 @@ async def handle_executed(request: web.Request) -> web.Response: async def broker_with_admin_middleware( admin_api_server: TestServer, ) -> AsyncGenerator[InMemoryBroker, None]: - broker = InMemoryBroker(await_inplace=True).with_middlewares( - TaskiqAdminMiddleware( - str(admin_api_server.make_url("/")), # URL тестового сервера - "supersecret", - taskiq_broker_name="InMemory", - ), + middleware = TaskiqAdminMiddleware( + str(admin_api_server.make_url("/")), # URL тестового сервера + "supersecret", + taskiq_broker_name="InMemory", ) + broker = InMemoryBroker(await_inplace=True).with_middlewares(middleware) broker.register_task(task_with_dataclass, task_name="task_with_dataclass") broker.register_task(task_with_typed_dict, task_name="task_with_typed_dict") @@ -55,6 +54,8 @@ async def broker_with_admin_middleware( await broker.startup() yield broker await broker.shutdown() + assert middleware._client is not None + assert middleware._client.closed async def task_with_dataclass(dto: DataclassDTO) -> None: diff --git a/tests/middlewares/admin_middleware/test_arguments_formatting.py b/tests/middlewares/admin_middleware/test_arguments_formatting.py index a65b5669..0d021625 100644 --- a/tests/middlewares/admin_middleware/test_arguments_formatting.py +++ b/tests/middlewares/admin_middleware/test_arguments_formatting.py @@ -22,7 +22,6 @@ class PydanticDTOFactory(ModelFactory[PydanticDTO]): __model__ = PydanticDTO -# @pytest.mark.skip class TestArgumentsFormattingInAdminMiddleware: @pytest.mark.parametrize( "dto_factory, task_name", diff --git a/tests/middlewares/test_simple_retry.py b/tests/middlewares/test_simple_retry.py index 783d98b7..1e7c7754 100644 --- a/tests/middlewares/test_simple_retry.py +++ b/tests/middlewares/test_simple_retry.py @@ -1,23 +1,11 @@ -import uuid -from unittest.mock import AsyncMock - -import pytest - -from taskiq.formatters.json_formatter import JSONFormatter from taskiq.message import TaskiqMessage from taskiq.middlewares.simple_retry_middleware import SimpleRetryMiddleware from taskiq.result import TaskiqResult +from tests.utils import AsyncQueueBroker -@pytest.fixture -def broker() -> AsyncMock: - mocked_broker = AsyncMock() - mocked_broker.id_generator = lambda: uuid.uuid4().hex - mocked_broker.formatter = JSONFormatter() - return mocked_broker - - -async def test_successful_retry(broker: AsyncMock) -> None: +async def test_successful_retry() -> None: + broker = AsyncQueueBroker() middleware = SimpleRetryMiddleware() middleware.set_broker(broker) await middleware.on_error( @@ -33,12 +21,13 @@ async def test_successful_retry(broker: AsyncMock) -> None: TaskiqResult(is_err=True, return_value=None, execution_time=0.0), Exception(), ) - resend: TaskiqMessage = broker.kick.await_args.args[0] + resend = broker.formatter.loads(broker.queue.get_nowait()) assert resend.task_name == "meme" assert resend.labels["_retries"] == "1" -async def test_no_retry(broker: AsyncMock) -> None: +async def test_no_retry() -> None: + broker = AsyncQueueBroker() middleware = SimpleRetryMiddleware() middleware.set_broker(broker) await middleware.on_error( @@ -52,10 +41,11 @@ async def test_no_retry(broker: AsyncMock) -> None: TaskiqResult(is_err=True, return_value=None, execution_time=0.0), Exception(), ) - broker.kick.assert_not_called() + assert broker.queue.empty() -async def test_max_retries(broker: AsyncMock) -> None: +async def test_max_retries() -> None: + broker = AsyncQueueBroker() middleware = SimpleRetryMiddleware(default_retry_count=3) middleware.set_broker(broker) await middleware.on_error( @@ -72,4 +62,4 @@ async def test_max_retries(broker: AsyncMock) -> None: TaskiqResult(is_err=True, return_value=None, execution_time=0.0), Exception(), ) - broker.kick.assert_not_called() + assert broker.queue.empty() From f1b2ab80e2bb570d9e17b2c6b19444700950b5d6 Mon Sep 17 00:00:00 2001 From: "sergei.romanchuk" Date: Mon, 13 Jul 2026 21:59:26 +0200 Subject: [PATCH 2/2] fix(tests): make OpenTelemetry span assertions order-independent Identify producer and consumer spans by SpanKind instead of relying on platform-dependent start-time ordering, fixing Windows CI failures. --- .../test_auto_instrumentation.py | 10 ++++++++-- tests/opentelemetry/test_tasks.py | 20 +++++++++++++++---- 2 files changed, 24 insertions(+), 6 deletions(-) diff --git a/tests/opentelemetry/test_auto_instrumentation.py b/tests/opentelemetry/test_auto_instrumentation.py index c7290869..f6350330 100644 --- a/tests/opentelemetry/test_auto_instrumentation.py +++ b/tests/opentelemetry/test_auto_instrumentation.py @@ -23,10 +23,16 @@ async def task_add(a: float, b: float) -> float: asyncio.run(test()) - spans = self.sorted_spans(self.memory_exporter.get_finished_spans()) + spans = self.memory_exporter.get_finished_spans() self.assertEqual(len(spans), 2) - consumer, producer = spans + spans_by_kind = {span.kind: span for span in spans} + self.assertEqual( + set(spans_by_kind), + {SpanKind.CONSUMER, SpanKind.PRODUCER}, + ) + consumer = spans_by_kind[SpanKind.CONSUMER] + producer = spans_by_kind[SpanKind.PRODUCER] self.assertEqual( consumer.name, diff --git a/tests/opentelemetry/test_tasks.py b/tests/opentelemetry/test_tasks.py index ab443f96..cd889de3 100644 --- a/tests/opentelemetry/test_tasks.py +++ b/tests/opentelemetry/test_tasks.py @@ -36,10 +36,16 @@ async def test() -> None: asyncio.run(test()) - spans = self.sorted_spans(self.memory_exporter.get_finished_spans()) + spans = self.memory_exporter.get_finished_spans() self.assertEqual(len(spans), 2) - consumer, producer = spans + spans_by_kind = {span.kind: span for span in spans} + self.assertEqual( + set(spans_by_kind), + {SpanKind.CONSUMER, SpanKind.PRODUCER}, + ) + consumer = spans_by_kind[SpanKind.CONSUMER] + producer = spans_by_kind[SpanKind.PRODUCER] self.assertEqual( consumer.name, @@ -85,10 +91,16 @@ async def test() -> None: asyncio.run(test()) - spans = self.sorted_spans(self.memory_exporter.get_finished_spans()) + spans = self.memory_exporter.get_finished_spans() self.assertEqual(len(spans), 2) - consumer, producer = spans + spans_by_kind = {span.kind: span for span in spans} + self.assertEqual( + set(spans_by_kind), + {SpanKind.CONSUMER, SpanKind.PRODUCER}, + ) + consumer = spans_by_kind[SpanKind.CONSUMER] + producer = spans_by_kind[SpanKind.PRODUCER] self.assertEqual( consumer.name,