feat: add tests to handling custom exc types

Author: GiperBoreipy

tests/middlewares/test_task_retry.py

@@ -2,7 +2,7 @@
 
 import pytest
 
-from taskiq import InMemoryBroker, SimpleRetryMiddleware
+from taskiq import InMemoryBroker, SimpleRetryMiddleware, SmartRetryMiddleware
 from taskiq.exceptions import NoResultError
 
 
@@ -151,3 +151,109 @@ def run_task() -> str:
 
     assert runs == 1
     assert str(resp.error) == str(runs)
+
+
+@pytest.mark.anyio
+async def test_retry_of_custom_exc_types_of_simple_middleware() -> None:
+    # test that the passed error will be handled
+    broker = InMemoryBroker().with_middlewares(
+        SimpleRetryMiddleware(
+            no_result_on_retry=True,
+            default_retry_label=True,
+            types_of_exceptions=(KeyError, ValueError),
+        )
+    )
+    runs = 0
+
+    @broker.task(max_retries=10)
+    def run_task() -> None:
+        nonlocal runs
+
+        runs += 1
+
+        raise ValueError()
+
+    task = await run_task.kiq()
+    resp = await task.wait_result(timeout=1)
+    with pytest.raises(ValueError):
+        resp.raise_for_error()
+
+    assert runs == 10
+
+    # test that an untransmitted error will not be handled
+    broker = InMemoryBroker().with_middlewares(
+        SimpleRetryMiddleware(
+            no_result_on_retry=True,
+            default_retry_label=True,
+            types_of_exceptions=(KeyError,),
+        )
+    )
+    runs = 0
+
+    @broker.task(max_retries=10)
+    def run_task2() -> None:
+        nonlocal runs
+
+        runs += 1
+
+        raise ValueError()
+
+    task = await run_task2.kiq()
+    resp = await task.wait_result(timeout=1)
+    with pytest.raises(ValueError):
+        resp.raise_for_error()
+
+    assert runs == 1
+
+
+@pytest.mark.anyio
+async def test_retry_of_custom_exc_types_of_smart_middleware() -> None:
+    # test that the passed error will be handled
+    broker = InMemoryBroker().with_middlewares(
+        SmartRetryMiddleware(
+            no_result_on_retry=True,
+            default_retry_label=True,
+            types_of_exceptions=(KeyError, ValueError),
+        )
+    )
+    runs = 0
+
+    @broker.task(max_retries=10)
+    def run_task() -> None:
+        nonlocal runs
+
+        runs += 1
+
+        raise ValueError()
+
+    task = await run_task.kiq()
+    resp = await task.wait_result(timeout=1)
+    with pytest.raises(ValueError):
+        resp.raise_for_error()
+
+    assert runs == 10
+
+    # test that an untransmitted error will not be handled
+    broker = InMemoryBroker().with_middlewares(
+        SmartRetryMiddleware(
+            no_result_on_retry=True,
+            default_retry_label=True,
+            types_of_exceptions=(KeyError,),
+        )
+    )
+    runs = 0
+
+    @broker.task(max_retries=10)
+    def run_task2() -> None:
+        nonlocal runs
+
+        runs += 1
+
+        raise ValueError()
+
+    task = await run_task2.kiq()
+    resp = await task.wait_result(timeout=1)
+    with pytest.raises(ValueError):
+        resp.raise_for_error()
+
+    assert runs == 1