feat: AbortPipeline error propagated to the last step task result of pipeline

Author: Юдыцкий Игорь

taskiq_pipelines/middleware.py

@@ -108,23 +108,24 @@ async def on_error(
             return
         if current_step_num == len(steps) - 1:
             return
-        await self.fail_pipeline(steps[-1].task_id)
+        await self.fail_pipeline(steps[-1].task_id, result.error)
 
-    async def fail_pipeline(self, last_task_id: str) -> None:
+    async def fail_pipeline(self, last_task_id: str, abort: AbortPipeline | None = None) -> None:
         """
         This function aborts pipeline.
 
         This is done by setting error result for
         the last task in the pipeline.
 
         :param last_task_id: id of the last task.
+        :param abort: caught earlier exception or default
         """
         await self.broker.result_backend.set_result(
             last_task_id,
             TaskiqResult(
                 is_err=True,
                 return_value=None,  # type: ignore
-                error=AbortPipeline("Execution aborted."),
+                error=abort or AbortPipeline("Execution aborted."),
                 execution_time=0,
                 log="Error found while executing pipeline.",
             ),

tests/test_steps.py

@@ -1,9 +1,10 @@
+import sys
 from typing import List
 
 import pytest
 from taskiq import InMemoryBroker
 
-from taskiq_pipelines import Pipeline, PipelineMiddleware
+from taskiq_pipelines import Pipeline, PipelineMiddleware, AbortPipeline
 
 
 @pytest.mark.anyio
@@ -42,3 +43,32 @@ def double(i: int) -> int:
     sent = await pipe.kiq(4)
     res = await sent.wait_result()
     assert res.return_value == list(map(double, ranger(4)))
+
+
+@pytest.mark.anyio
+async def test_abort_pipeline() -> None:
+    """Test AbortPipeline."""
+    broker = InMemoryBroker().with_middlewares(PipelineMiddleware())
+    text = 'task was aborted'
+
+    @broker.task
+    def normal_task(i: bool) -> bool:
+        return i
+
+    @broker.task
+    def aborting_task(i: int) -> bool:
+        if i:
+            raise AbortPipeline(text)
+        return True
+
+    pipe = Pipeline(broker, aborting_task).call_next(normal_task)
+    sent = await pipe.kiq(0)
+    res = await sent.wait_result()
+    assert res.is_err is False
+    assert res.return_value is True
+    assert res.error is None
+    sent = await pipe.kiq(1)
+    res = await sent.wait_result()
+    assert res.is_err is True
+    assert res.return_value is None
+    assert res.error.args[0] == text