Keep retry exception filters reusable

Author: IlyaKhalafi

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,

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(

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()