fix: add fake task and kicker

Author: Lancetnik

taskiq_faststream/broker.py

@@ -8,8 +8,8 @@
 from faststream.types import SendableMessage
 from taskiq import AsyncBroker, BrokerMessage
 from taskiq.acks import AckableMessage
-from taskiq.decor import AsyncTaskiqDecoratedTask
 
+from taskiq_faststream.task import PatchedTaskiqDecoratedTask
 from taskiq_faststream.types import ScheduledTask
 from taskiq_faststream.utils import resolve_msg
 
@@ -33,6 +33,7 @@ class BrokerWrapper(AsyncBroker):
     def __init__(self, broker: BrokerAsyncUsecase[typing.Any, typing.Any]) -> None:
         super().__init__()
         self.broker = broker
+        self.decorator_class = PatchedTaskiqDecoratedTask
 
     async def startup(self) -> None:
         """Startup wrapped FastStream broker."""
@@ -76,7 +77,7 @@ def task(  # type: ignore[override]
         *,
         schedule: typing.List[ScheduledTask],
         **kwargs: PublishParameters,
-    ) -> "AsyncTaskiqDecoratedTask[[], None]":
+    ) -> "PatchedTaskiqDecoratedTask[[], None]":
         """Register FastStream scheduled task.
 
         Args:
@@ -108,6 +109,7 @@ class AppWrapper(BrokerWrapper):
     def __init__(self, app: FastStream) -> None:
         super(BrokerWrapper, self).__init__()
         self.app = app
+        self.decorator_class = PatchedTaskiqDecoratedTask
 
     async def startup(self) -> None:
         """Startup wrapped FastStream."""

taskiq_faststream/kicker.py

@@ -0,0 +1,13 @@
+from typing import Any
+
+from taskiq.kicker import AsyncKicker, _FuncParams, _ReturnType
+from taskiq.message import TaskiqMessage
+
+
+class LabelRespectKicker(AsyncKicker[_FuncParams, _ReturnType]):
+    """Patched kicker doesn't cast labels to str."""
+
+    def _prepare_message(self, *args: Any, **kwargs: Any) -> TaskiqMessage:
+        msg = super()._prepare_message(*args, **kwargs)
+        msg.labels = self.labels
+        return msg

taskiq_faststream/task.py

@@ -0,0 +1,23 @@
+from taskiq.decor import AsyncTaskiqDecoratedTask as Task
+from taskiq.decor import _FuncParams, _ReturnType
+
+from taskiq_faststream.kicker import LabelRespectKicker
+
+
+class PatchedTaskiqDecoratedTask(Task[_FuncParams, _ReturnType]):
+    """Patched Decorated Task has a patched kicker."""
+
+    def kicker(self) -> LabelRespectKicker[_FuncParams, _ReturnType]:
+        """
+        This function returns kicker object.
+
+        Kicker is a object that can modify kiq request
+        before sending it.
+
+        :return: AsyncKicker instance.
+        """
+        return LabelRespectKicker(
+            task_name=self.task_name,
+            broker=self.broker,
+            labels=self.labels,
+        )