Fixed middleware testing.

Author: s3rius

taskiq/decor.py

@@ -1,3 +1,4 @@
+import copy
 import sys
 from collections.abc import Coroutine
 from datetime import datetime
@@ -8,13 +9,15 @@
     Callable,
     Dict,
     Generic,
+    Optional,
     TypeVar,
     Union,
     overload,
 )
 
 from typing_extensions import ParamSpec
 
+from taskiq.abc.middleware import TaskiqMiddleware
 from taskiq.kicker import AsyncKicker
 from taskiq.scheduler.created_schedule import CreatedSchedule
 from taskiq.task import AsyncTaskiqTask
@@ -51,11 +54,15 @@ def __init__(
         task_name: str,
         original_func: Callable[_FuncParams, _ReturnType],
         labels: Dict[str, Any],
+        extra_middlewares: Optional[list[TaskiqMiddleware]] = None,
     ) -> None:
         self.broker = broker
         self.task_name = task_name
         self.original_func = original_func
         self.labels = labels
+        self.middlewares = copy.copy(broker.middlewares)
+        if extra_middlewares:
+            self.middlewares.extend(extra_middlewares)
 
         # This is a hack to make ProcessPoolExecutor work
         # with decorated functions.

taskiq/kicker.py

@@ -153,7 +153,7 @@ async def kiq(
         logger.debug(
             f"Kicking {self.task_name} with args={args} and kwargs={kwargs}.",
         )
-        message = self._prepare_message(*args, **kwargs)
+        message = self.get_message(*args, **kwargs)
         for middleware in self.broker.middlewares:
             if middleware.__class__.pre_send != TaskiqMiddleware.pre_send:
                 message = await maybe_awaitable(middleware.pre_send(message))
@@ -191,7 +191,7 @@ async def schedule_by_cron(
         schedule_id = self.custom_schedule_id
         if schedule_id is None:
             schedule_id = self.broker.id_generator()
-        message = self._prepare_message(*args, **kwargs)
+        message = self.get_message(*args, **kwargs)
         cron_offset = None
         if isinstance(cron, CronSpec):
             cron_str = cron.to_cron()
@@ -228,7 +228,7 @@ async def schedule_by_time(
         schedule_id = self.custom_schedule_id
         if schedule_id is None:
             schedule_id = self.broker.id_generator()
-        message = self._prepare_message(*args, **kwargs)
+        message = self.get_message(*args, **kwargs)
         scheduled = ScheduledTask(
             schedule_id=schedule_id,
             task_name=message.task_name,
@@ -261,10 +261,10 @@ def _prepare_arg(cls, arg: Any) -> Any:
             arg = asdict(arg)
         return arg
 
-    def _prepare_message(
+    def get_message(
         self,
-        *args: Any,
-        **kwargs: Any,
+        *args: _FuncParams.args,
+        **kwargs: _FuncParams.kwargs,
     ) -> TaskiqMessage:
         """
         Create a message from args and kwargs.

taskiq/receiver/receiver.py

@@ -125,7 +125,7 @@ async def callback(  # noqa: C901, PLR0912
             "Function for task %s is resolved. Executing...",
             taskiq_msg.task_name,
         )
-        for middleware in self.broker.middlewares:
+        for middleware in task.middlewares:
             if middleware.__class__.pre_execute != TaskiqMiddleware.pre_execute:
                 taskiq_msg = await maybe_awaitable(
                     middleware.pre_execute(
@@ -150,21 +150,32 @@ async def callback(  # noqa: C901, PLR0912
             message=taskiq_msg,
         )
 
+        if result.is_err is not None:
+            for middleware in task.middlewares:
+                if middleware.__class__.on_error != TaskiqMiddleware.on_error:
+                    await maybe_awaitable(
+                        middleware.on_error(
+                            taskiq_msg,
+                            result,
+                            result.error,  # type: ignore
+                        ),
+                    )
+
         if self.ack_time == AcknowledgeType.WHEN_EXECUTED and isinstance(
             message,
             AckableMessage,
         ):
             await maybe_awaitable(message.ack())
 
-        for middleware in self.broker.middlewares:
+        for middleware in task.middlewares:
             if middleware.__class__.post_execute != TaskiqMiddleware.post_execute:
                 await maybe_awaitable(middleware.post_execute(taskiq_msg, result))
 
         try:
             if not isinstance(result.error, NoResultError):
                 await self.broker.result_backend.set_result(taskiq_msg.task_id, result)
 
-                for middleware in self.broker.middlewares:
+                for middleware in task.middlewares:
                     if middleware.__class__.post_save != TaskiqMiddleware.post_save:
                         await maybe_awaitable(middleware.post_save(taskiq_msg, result))
 
@@ -183,7 +194,7 @@ async def callback(  # noqa: C901, PLR0912
         ):
             await maybe_awaitable(message.ack())
 
-    async def run_task(  # noqa: C901, PLR0912, PLR0915
+    async def run_task(  # noqa: C901
         self,
         target: Callable[..., Any],
         message: TaskiqMessage,
@@ -304,17 +315,6 @@ async def run_task(  # noqa: C901, PLR0912, PLR0915
             error=found_exception,
             labels=message.labels,
         )
-        # If exception is found we execute middlewares.
-        if found_exception is not None:
-            for middleware in self.broker.middlewares:
-                if middleware.__class__.on_error != TaskiqMiddleware.on_error:
-                    await maybe_awaitable(
-                        middleware.on_error(
-                            message,
-                            result,
-                            found_exception,
-                        ),
-                    )
 
         return result
 

taskiq/task_creator.py

@@ -15,6 +15,7 @@
 
 from typing_extensions import ParamSpec, Self
 
+from taskiq.abc.middleware import TaskiqMiddleware
 from taskiq.decor import AsyncTaskiqDecoratedTask
 from taskiq.utils import remove_suffix
 
@@ -37,6 +38,7 @@ def __init__(self, broker: "AsyncBroker") -> None:
         self._broker = broker
         self._task_name: Optional[str] = None
         self._labels: Dict[str, Any] = {}
+        self._middlewares: list[TaskiqMiddleware] = []
 
     def name(self, name: str) -> Self:
         """Assign custom name to the task."""
@@ -48,6 +50,11 @@ def labels(self, **labels: Any) -> Self:
         self._labels = labels
         return self
 
+    def middlewares(self, *middlewares: TaskiqMiddleware) -> Self:
+        """Assign custom middlewares to the task."""
+        self._middlewares = list(middlewares)
+        return self
+
     def make_task(
         self,
         task_name: str,
@@ -59,8 +66,21 @@ def make_task(
             original_func=func,
             labels=self._labels,
             task_name=task_name,
+            extra_middlewares=self._middlewares,
         )
 
+    def __resolve_name(self, func: Callable[..., Any]) -> str:
+        """Resolve name of the function."""
+        fmodule = func.__module__
+        if fmodule == "__main__":  # pragma: no cover
+            fmodule = ".".join(
+                remove_suffix(sys.argv[0], ".py").split(os.path.sep),
+            )
+        fname = func.__name__
+        if fname == "<lambda>":
+            fname = f"lambda_{uuid4().hex}"
+        return f"{fmodule}:{fname}"
+
     @overload
     def __call__(
         self,
@@ -131,15 +151,7 @@ def inner(
             ) -> AsyncTaskiqDecoratedTask[_FuncParams, _ReturnType]:
                 inner_task_name = self._task_name
                 if inner_task_name is None:
-                    fmodule = func.__module__
-                    if fmodule == "__main__":  # pragma: no cover
-                        fmodule = ".".join(
-                            remove_suffix(sys.argv[0], ".py").split(os.path.sep),
-                        )
-                    fname = func.__name__
-                    if fname == "<lambda>":
-                        fname = f"lambda_{uuid4().hex}"
-                    inner_task_name = f"{fmodule}:{fname}"
+                    inner_task_name = self.__resolve_name(func)
 
                 wrapper = wraps(func)
                 decorated_task = wrapper(

tests/abc/test_broker.py

@@ -40,7 +40,7 @@ def test_decorator_with_name_success() -> None:
     """Test that task_name is successfully set."""
     tbrok = _TestBroker()
 
-    @tbrok.task(task_name="my_task")
+    @tbrok.task.name("my_task")
     async def test_func() -> None:
         """Some test function."""
 
@@ -52,7 +52,7 @@ def test_decorator_with_labels_success() -> None:
     """Tests that labels are assigned for task as is."""
     tbrok = _TestBroker()
 
-    @tbrok.task(label1=1, label2=2)
+    @tbrok.task.labels(label1=1, label2=2)
     async def test_func() -> None:
         """Some test function."""
 

tests/api/test_scheduler.py

@@ -1,6 +1,6 @@
 import asyncio
 import contextlib
-from datetime import datetime, timedelta
+from datetime import datetime, timedelta, timezone
 
 import pytest
 
@@ -16,7 +16,9 @@ async def test_successful() -> None:
     scheduler = TaskiqScheduler(broker, sources=[LabelScheduleSource(broker)])
     scheduler_task = asyncio.create_task(run_scheduler_task(scheduler))
 
-    @broker.task(schedule=[{"time": datetime.utcnow() - timedelta(seconds=1)}])
+    @broker.task.labels(
+        schedule=[{"time": datetime.now(timezone.utc) - timedelta(seconds=1)}],
+    )
     def _() -> None:
         ...
 
@@ -31,7 +33,7 @@ async def test_cancelation() -> None:
     broker = AsyncQueueBroker()
     scheduler = TaskiqScheduler(broker, sources=[LabelScheduleSource(broker)])
 
-    @broker.task(schedule=[{"time": datetime.utcnow()}])
+    @broker.task.labels(schedule=[{"time": datetime.now(timezone.utc)}])
     def _() -> None:
         ...
 

tests/middlewares/test_task_retry.py

@@ -14,7 +14,7 @@ async def test_wait_result() -> None:
     )
     runs = 0
 
-    @broker.task(retry_on_error=True)
+    @broker.task.labels(retry_on_error=True)
     async def run_task() -> str:
         nonlocal runs
 
@@ -40,7 +40,7 @@ async def test_wait_result_error() -> None:
     runs = 0
     lock = asyncio.Lock()
 
-    @broker.task(retry_on_error=True)
+    @broker.task.labels(retry_on_error=True)
     async def run_task() -> str:
         nonlocal runs, lock
 
@@ -74,7 +74,7 @@ async def test_wait_result_no_result() -> None:
     runs = 0
     lock = asyncio.Lock()
 
-    @broker.task(retry_on_error=True)
+    @broker.task.labels(retry_on_error=True)
     async def run_task() -> str:
         nonlocal runs, done, lock
 
@@ -111,7 +111,7 @@ async def test_max_retries() -> None:
     )
     runs = 0
 
-    @broker.task(max_retries=10)
+    @broker.task.labels(max_retries=10)
     def run_task() -> str:
         nonlocal runs
 
@@ -137,7 +137,7 @@ async def test_no_retry() -> None:
     )
     runs = 0
 
-    @broker.task(retry_on_error=False, max_retries=10)
+    @broker.task.labels(retry_on_error=False, max_retries=10)
     def run_task() -> str:
         nonlocal runs
 

tests/receiver/test_receiver.py

@@ -2,18 +2,16 @@
 import random
 import time
 from concurrent.futures import ThreadPoolExecutor
-from typing import Any, ClassVar, List, Optional
+from typing import Optional
 
 import pytest
 from taskiq_dependencies import Depends
 
 from taskiq.abc.broker import AckableMessage, AsyncBroker
-from taskiq.abc.middleware import TaskiqMiddleware
 from taskiq.brokers.inmemory_broker import InMemoryBroker
 from taskiq.exceptions import NoResultError, TaskiqResultTimeoutError
 from taskiq.message import TaskiqMessage
 from taskiq.receiver import Receiver
-from taskiq.result import TaskiqResult
 from tests.utils import AsyncQueueBroker
 
 
@@ -152,43 +150,6 @@ def test_func() -> None:
     assert result.is_err
 
 
-@pytest.mark.anyio
-async def test_run_task_exception_middlewares() -> None:
-    """Tests that run_task can run sync tasks."""
-
-    class _TestMiddleware(TaskiqMiddleware):
-        found_exceptions: ClassVar[List[BaseException]] = []
-
-        def on_error(
-            self,
-            message: "TaskiqMessage",
-            result: "TaskiqResult[Any]",
-            exception: BaseException,
-        ) -> None:
-            self.found_exceptions.append(exception)
-
-    def test_func() -> None:
-        raise ValueError
-
-    broker = InMemoryBroker().with_middlewares(_TestMiddleware())
-    receiver = get_receiver(broker)
-
-    result = await receiver.run_task(
-        test_func,
-        TaskiqMessage(
-            task_id="",
-            task_name="",
-            labels={},
-            args=[],
-            kwargs={},
-        ),
-    )
-    assert result.return_value is None
-    assert result.is_err
-    assert len(_TestMiddleware.found_exceptions) == 1
-    assert _TestMiddleware.found_exceptions[0].__class__ is ValueError
-
-
 @pytest.mark.anyio
 async def test_callback_success() -> None:
     """Test that callback function works well."""