From c88ffa213c708c7726ee66a4fd5c60744dce3abf Mon Sep 17 00:00:00 2001 From: IlyaKhalafi Date: Mon, 19 Oct 2026 21:14:32 +0000 Subject: [PATCH] Keep retry exception filters reusable --- taskiq/middlewares/simple_retry_middleware.py | 4 +- taskiq/middlewares/smart_retry_middleware.py | 4 +- tests/middlewares/test_retry_iterables.py | 105 ++++++++++++++++++ 3 files changed, 111 insertions(+), 2 deletions(-) create mode 100644 tests/middlewares/test_retry_iterables.py diff --git a/taskiq/middlewares/simple_retry_middleware.py b/taskiq/middlewares/simple_retry_middleware.py index 4c2f9a45..9d31a231 100644 --- a/taskiq/middlewares/simple_retry_middleware.py +++ b/taskiq/middlewares/simple_retry_middleware.py @@ -24,7 +24,9 @@ def __init__( self.default_retry_count = default_retry_count self.default_retry_label = default_retry_label self.no_result_on_retry = no_result_on_retry - self.types_of_exceptions = types_of_exceptions + self.types_of_exceptions = ( + tuple(types_of_exceptions) if types_of_exceptions is not None else None + ) def _get_types_of_exceptions( self, diff --git a/taskiq/middlewares/smart_retry_middleware.py b/taskiq/middlewares/smart_retry_middleware.py index 9752b07f..46ea8133 100644 --- a/taskiq/middlewares/smart_retry_middleware.py +++ b/taskiq/middlewares/smart_retry_middleware.py @@ -61,7 +61,9 @@ def __init__( self.use_delay_exponent = use_delay_exponent self.max_delay_exponent = max_delay_exponent self.schedule_source = schedule_source - self.types_of_exceptions = types_of_exceptions + self.types_of_exceptions = ( + tuple(types_of_exceptions) if types_of_exceptions is not None else None + ) if not isinstance(schedule_source, (ScheduleSource, type(None))): raise TypeError( diff --git a/tests/middlewares/test_retry_iterables.py b/tests/middlewares/test_retry_iterables.py new file mode 100644 index 00000000..90534a87 --- /dev/null +++ b/tests/middlewares/test_retry_iterables.py @@ -0,0 +1,105 @@ +from collections.abc import Iterable +from typing import Any + +import pytest + +from taskiq import ( + InMemoryBroker, + SimpleRetryMiddleware, + SmartRetryMiddleware, + TaskiqMessage, + TaskiqResult, +) +from taskiq.message import BrokerMessage + + +class RecordingBroker(InMemoryBroker): + def __init__(self) -> None: + super().__init__() + self.sent: list[BrokerMessage] = [] + + async def kick(self, message: BrokerMessage) -> None: + self.sent.append(message) + + +@pytest.mark.parametrize( + "middleware_class", + [SimpleRetryMiddleware, SmartRetryMiddleware], +) +@pytest.mark.parametrize("filter_kind", ["generator", "list", "empty", "none"]) +async def test_exception_filter_is_reusable( + middleware_class: type[SimpleRetryMiddleware] | type[SmartRetryMiddleware], + filter_kind: str, +) -> None: + exception_types: Iterable[type[BaseException]] | None + if filter_kind == "generator": + exception_types = (kind for kind in [ValueError]) + elif filter_kind == "list": + exception_types = [ValueError] + elif filter_kind == "empty": + exception_types = iter([]) + else: + exception_types = None + + broker = RecordingBroker() + middleware = middleware_class( + default_retry_label=True, + types_of_exceptions=exception_types, + ) + broker.with_middlewares(middleware) + try: + for task_id in ["first", "second"]: + error = ValueError("retry") + message = TaskiqMessage( + task_id=task_id, + task_name="test", + labels={}, + args=[], + kwargs={}, + ) + result: TaskiqResult[Any] = TaskiqResult( + is_err=True, + log=None, + return_value=None, + execution_time=0, + error=error, + ) + await middleware.on_error(message, result, error) + assert len(broker.sent) == (0 if filter_kind == "empty" else 2) + finally: + await broker.shutdown() + + +@pytest.mark.parametrize( + "middleware_class", + [SimpleRetryMiddleware, SmartRetryMiddleware], +) +async def test_nonmatching_error_does_not_consume_filter( + middleware_class: type[SimpleRetryMiddleware] | type[SmartRetryMiddleware], +) -> None: + broker = RecordingBroker() + middleware = middleware_class( + default_retry_label=True, + types_of_exceptions=(kind for kind in [ValueError]), + ) + broker.with_middlewares(middleware) + try: + for error in [TypeError("skip"), ValueError("retry")]: + message = TaskiqMessage( + task_id=str(type(error)), + task_name="test", + labels={}, + args=[], + kwargs={}, + ) + result: TaskiqResult[Any] = TaskiqResult( + is_err=True, + log=None, + return_value=None, + execution_time=0, + error=error, + ) + await middleware.on_error(message, result, error) + assert len(broker.sent) == 1 + finally: + await broker.shutdown()