fix: retry_middleware (#126)

Co-authored-by: Anton <[email protected]>

Author: Sobes76rus

taskiq/middlewares/retry_middleware.py

@@ -3,6 +3,7 @@
 from typing import Any
 
 from taskiq.abc.middleware import TaskiqMiddleware
+from taskiq.exceptions import NoResultError
 from taskiq.message import TaskiqMessage
 from taskiq.result import TaskiqResult
 
@@ -15,8 +16,12 @@ class SimpleRetryMiddleware(TaskiqMiddleware):
     def __init__(
         self,
         default_retry_count: int = 3,
+        default_retry_label: bool = False,
+        no_result_on_retry: bool = True,
     ) -> None:
         self.default_retry_count = default_retry_count
+        self.default_retry_label = default_retry_label
+        self.no_result_on_retry = no_result_on_retry
 
     async def on_error(
         self,
@@ -37,24 +42,34 @@ async def on_error(
         :param result: execution result.
         :param exception: found exception.
         """
+        # Valid exception
+        if isinstance(exception, NoResultError):
+            return
+
         retry_on_error = message.labels.get("retry_on_error")
+        if retry_on_error is None:
+            retry_on_error = "true" if self.default_retry_label else "false"
         # Check if retrying is enabled for the task.
-        if retry_on_error is None or retry_on_error.lower() != "true":
+        if retry_on_error.lower() != "true":
             return
         new_msg = deepcopy(message)
+
         # Getting number of previous retries.
         retries = int(new_msg.labels.get("_retries", 0)) + 1
         new_msg.labels["_retries"] = str(retries)
         max_retries = int(new_msg.labels.get("max_retries", self.default_retry_count))
+
         if retries < max_retries:
             logger.info(
                 "Task '%s' invocation failed. Retrying.",
                 message.task_name,
             )
-            new_msg.labels["_parent"] = message.task_id
-            new_msg.task_id = self.broker.id_generator()
             broker_message = self.broker.formatter.dumps(message=new_msg)
             await self.broker.kick(broker_message)
+
+            if self.no_result_on_retry:
+                result.error = NoResultError()
+
         else:
             logger.warning(
                 "Task '%s' invocation failed. Maximum retries count is reached.",

tests/middlewares/test_simple_retry.py

@@ -37,7 +37,6 @@ async def test_successful_retry(broker: AsyncMock) -> None:
     resend: TaskiqMessage = broker.kick.await_args.args[0]
     assert resend.task_name == "meme"
     assert resend.labels["_retries"] == "1"
-    assert resend.labels["_parent"] == "test_id"
 
 
 @pytest.mark.anyio

tests/middlewares/test_task_retry.py

@@ -0,0 +1,152 @@
+import asyncio
+import time
+
+import pytest
+
+from taskiq import InMemoryBroker, SimpleRetryMiddleware
+from taskiq.exceptions import NoResultError
+
+
+@pytest.mark.anyio
+async def test_wait_result() -> None:
+    """Tests wait_result."""
+
+    broker = InMemoryBroker().with_middlewares(
+        SimpleRetryMiddleware(no_result_on_retry=True),
+    )
+    runs = 0
+
+    @broker.task(retry_on_error=True)
+    def run_task() -> str:
+        nonlocal runs  # noqa: WPS420
+
+        if runs == 0:
+            runs += 1
+            raise Exception("Retry")
+
+        time.sleep(0.2)
+        return "hello world!"
+
+    task = await run_task.kiq()
+    resp = await task.wait_result(0.1, timeout=1)
+
+    assert resp.return_value == "hello world!"
+
+
+@pytest.mark.anyio
+async def test_wait_result_error() -> None:
+    """Tests wait_result."""
+
+    broker = InMemoryBroker().with_middlewares(
+        SimpleRetryMiddleware(no_result_on_retry=False),
+    )
+    runs = 0
+
+    @broker.task(retry_on_error=True)
+    def run_task() -> str:
+        nonlocal runs  # noqa: WPS420
+
+        if runs == 0:
+            runs += 1
+            raise ValueError("Retry")
+
+        time.sleep(0.2)
+        return "hello world!"
+
+    task = await run_task.kiq()
+    resp = await task.wait_result(0.1, timeout=1)
+    with pytest.raises(ValueError):
+        resp.raise_for_error()
+
+    await asyncio.sleep(0.2)
+    resp = await task.wait_result(timeout=1)
+    assert resp.return_value == "hello world!"
+
+
+@pytest.mark.anyio
+async def test_wait_result_no_result() -> None:
+    """Tests wait_result."""
+
+    broker = InMemoryBroker().with_middlewares(
+        SimpleRetryMiddleware(no_result_on_retry=False),
+    )
+    done = False
+    runs = 0
+
+    @broker.task(retry_on_error=True)
+    def run_task() -> str:
+        nonlocal runs, done  # noqa: WPS420
+
+        if runs == 0:
+            runs += 1
+            raise ValueError("Retry")
+
+        time.sleep(0.2)
+        done = True
+        raise NoResultError()
+
+    task = await run_task.kiq()
+    resp = await task.wait_result(0.1, timeout=1)
+    with pytest.raises(ValueError):
+        resp.raise_for_error()
+
+    await asyncio.sleep(0.2)
+    resp = await task.wait_result(timeout=1)
+    with pytest.raises(ValueError):
+        resp.raise_for_error()
+
+    assert done
+
+
+@pytest.mark.anyio
+async def test_max_retries() -> None:
+    """Tests wait_result."""
+
+    broker = InMemoryBroker().with_middlewares(
+        SimpleRetryMiddleware(
+            no_result_on_retry=True,
+            default_retry_label=True,
+        ),
+    )
+    runs = 0
+
+    @broker.task(max_retries=10)
+    def run_task() -> str:
+        nonlocal runs  # noqa: WPS420
+
+        runs += 1
+        raise ValueError(runs)
+
+    task = await run_task.kiq()
+    resp = await task.wait_result(timeout=1)
+    with pytest.raises(ValueError):
+        resp.raise_for_error()
+
+    assert runs == 10
+    assert str(resp.error) == str(runs)
+
+
+@pytest.mark.anyio
+async def test_no_retry() -> None:
+    broker = InMemoryBroker().with_middlewares(
+        SimpleRetryMiddleware(
+            no_result_on_retry=True,
+            default_retry_label=True,
+        ),
+    )
+    runs = 0
+
+    @broker.task(retry_on_error=False, max_retries=10)
+    def run_task() -> str:
+        nonlocal runs  # noqa: WPS420
+
+        runs += 1
+        raise ValueError(runs)
+
+    task = await run_task.kiq()
+    resp = await task.wait_result(timeout=1)
+    with pytest.raises(ValueError):
+        resp.raise_for_error()
+
+    assert runs == 1
+    assert str(resp.error) == str(runs)