feat: add BrokerWrapper

Author: Lancetnik

pyproject.toml

@@ -176,7 +176,7 @@ exclude = [".venv/"]
 ]
 
 [tool.ruff.pydocstyle]
-convention = "pep257"
+convention = "google"
 ignore-decorators = ["typing.overload"]
 
 [tool.ruff.pylint]

taskiq_faststream/__init__.py

@@ -1 +1,3 @@
-"""Project was generated using https://github.com/taskiq-python/project-template/."""
+from taskiq_faststream.broker import BrokerWrapper
+
+__all__ = ("BrokerWrapper",)

taskiq_faststream/broker.py

@@ -0,0 +1,92 @@
+import typing
+import warnings
+
+from faststream._compat import TypeAlias, override
+from faststream.broker.core.asyncronous import BrokerAsyncUsecase
+from faststream.types import SendableMessage
+from taskiq import AsyncBroker, BrokerMessage
+from taskiq.acks import AckableMessage
+from taskiq.decor import AsyncTaskiqDecoratedTask
+
+from taskiq_faststream.types import ScheduledTask
+from taskiq_faststream.utils import resolve_msg
+
+PublishParameters: TypeAlias = typing.Any
+
+
+class BrokerWrapper(AsyncBroker):
+    """Wrap FastStream broker to taskiq compatible object.
+
+    Attributes:
+        broker : FastStream wrapped broker.
+
+    Methods:
+        __init__ : Initializes the object.
+        startup : Startup wrapper FastStream broker.
+        shutdown : Shutdown wrapper FastStream broker.
+        kick : Call wrapped FastStream broker `publish` method.
+        task : Register FastStream scheduled task.
+    """
+
+    def __init__(self, broker: BrokerAsyncUsecase[typing.Any, typing.Any]) -> None:
+        super().__init__()
+        self.broker = broker
+
+    async def startup(self) -> None:
+        """Startup wrapper FastStream broker."""
+        await super().startup()
+        await self.broker.start()
+
+    async def shutdown(self) -> None:
+        """Shutdown wrapper FastStream broker."""
+        await self.broker.close()
+        await super().shutdown()
+
+    async def kick(self, message: BrokerMessage) -> None:
+        """Call wrapped FastStream broker `publish` method."""
+        labels = message.labels
+        labels.pop("schedule", None)
+        msg = await resolve_msg(labels.pop("message", message.message))
+        await self.broker.publish(msg, **labels)
+
+    async def listen(
+        self,
+    ) -> typing.AsyncGenerator[typing.Union[bytes, AckableMessage], None]:
+        """Not supported method."""
+        warnings.warn(
+            message=(
+                f"{self.__class__.__name__} doesn't support `listen` method. "
+                "Please, use it only to register a task."
+            ),
+            category=RuntimeWarning,
+            stacklevel=1,
+        )
+
+        while True:
+            yield b""
+
+    @override
+    def task(  # type: ignore[override]
+        self,
+        message: typing.Union[
+            None,
+            SendableMessage,
+            typing.Callable[[], SendableMessage],
+            typing.Callable[[], typing.Awaitable[SendableMessage]],
+        ] = None,
+        *,
+        schedule: typing.List[ScheduledTask],
+        **kwargs: PublishParameters,
+    ) -> AsyncTaskiqDecoratedTask[[], None]:
+        """Register FastStream scheduled task.
+
+        Args:
+            message: object to send or sync/async message generation callback.
+            schedule: scheduler parameters list.
+            kwargs: `broker.publish(...)` arguments.
+        """
+        return super().task(
+            message=message,
+            schedule=schedule,
+            **kwargs,
+        )(lambda: None)

taskiq_faststream/types.py

@@ -0,0 +1,15 @@
+from datetime import datetime, timedelta
+from typing import Optional, Union
+
+from faststream._compat import TypedDict
+
+
+class ScheduledTask(TypedDict, total=False):
+    """Store information about scheduled tasks.
+
+    https://taskiq-python.github.io/available-components/schedule-sources.html
+    """
+
+    cron: str
+    cron_offset: Union[str, timedelta, None]
+    time: Optional[datetime]

taskiq_faststream/utils.py

@@ -0,0 +1,29 @@
+import typing
+
+from faststream.types import SendableMessage
+from faststream.utils.functions import to_async
+
+
+async def resolve_msg(
+    msg: typing.Union[
+        None,
+        SendableMessage,
+        typing.Callable[[], SendableMessage],
+        typing.Callable[[], typing.Awaitable[SendableMessage]],
+    ],
+) -> SendableMessage:
+    """Resolve message generation callback.
+
+    Args:
+        msg: object to send or sync/async message generation callback.
+
+    Returns:
+        The message to send
+    """
+    if callable(msg):
+        get_msg = typing.cast(
+            typing.Callable[[], typing.Awaitable[SendableMessage]],
+            to_async(msg),
+        )
+        msg = await get_msg()
+    return msg

tests/conftest.py

@@ -1,4 +1,8 @@
+from unittest.mock import MagicMock
+from uuid import uuid4
+
 import pytest
+from anyio._backends._asyncio import Event
 
 
 @pytest.fixture(scope="session")
@@ -10,3 +14,18 @@ def anyio_backend() -> str:
     :return: backend name.
     """
     return "asyncio"
+
+
+@pytest.fixture()
+def subject() -> str:
+    return uuid4().hex
+
+
+@pytest.fixture()
+def mock() -> MagicMock:
+    return MagicMock()
+
+
+@pytest.fixture()
+def event() -> Event:
+    return Event()

tests/test_broker.py

@@ -0,0 +1,12 @@
+import pytest
+
+from taskiq_faststream import BrokerWrapper
+
+
+@pytest.mark.anyio
+async def test_warning() -> None:
+    broker = BrokerWrapper(None)
+
+    with pytest.warns(RuntimeWarning):
+        async for _ in broker.listen():
+            break

tests/test_kafka.py

@@ -0,0 +1,63 @@
+import asyncio
+from datetime import datetime
+from unittest.mock import MagicMock
+
+import pytest
+from anyio import Event
+from faststream.kafka import KafkaBroker, TestKafkaBroker
+from faststream.utils.functions import timeout_scope
+from taskiq import TaskiqScheduler
+from taskiq.cli.scheduler.args import SchedulerArgs
+from taskiq.cli.scheduler.run import run_scheduler
+from taskiq.schedule_sources import LabelScheduleSource
+
+from taskiq_faststream import BrokerWrapper
+
+
+@pytest.fixture
+def broker() -> KafkaBroker:
+    return KafkaBroker()
+
+
+@pytest.mark.anyio
+async def test_time_task(
+    subject: str,
+    broker: KafkaBroker,
+    mock: MagicMock,
+    event: Event,
+) -> None:
+    @broker.subscriber(subject)
+    async def handler(msg: str) -> None:
+        event.set()
+        mock(msg)
+
+    taskiq_broker = BrokerWrapper(broker)
+
+    taskiq_broker.task(
+        "Hi!",
+        topic=subject,
+        schedule=[
+            {
+                "time": datetime.utcnow(),
+            },
+        ],
+    )
+
+    async with TestKafkaBroker(broker):
+        task = asyncio.create_task(
+            run_scheduler(
+                SchedulerArgs(
+                    scheduler=TaskiqScheduler(
+                        broker=taskiq_broker,
+                        sources=[LabelScheduleSource(taskiq_broker)],
+                    ),
+                    modules=[],
+                ),
+            ),
+        )
+
+        with timeout_scope(3.0, True):
+            await event.wait()
+
+    mock.assert_called_once_with("Hi!")
+    task.cancel()

tests/test_nats.py

@@ -0,0 +1,63 @@
+import asyncio
+from datetime import datetime
+from unittest.mock import MagicMock
+
+import pytest
+from anyio import Event
+from faststream.nats import NatsBroker, TestNatsBroker
+from faststream.utils.functions import timeout_scope
+from taskiq import TaskiqScheduler
+from taskiq.cli.scheduler.args import SchedulerArgs
+from taskiq.cli.scheduler.run import run_scheduler
+from taskiq.schedule_sources import LabelScheduleSource
+
+from taskiq_faststream import BrokerWrapper
+
+
+@pytest.fixture
+def broker() -> NatsBroker:
+    return NatsBroker()
+
+
+@pytest.mark.anyio
+async def test_time_task(
+    subject: str,
+    broker: NatsBroker,
+    mock: MagicMock,
+    event: Event,
+) -> None:
+    @broker.subscriber(subject)
+    async def handler(msg: str) -> None:
+        event.set()
+        mock(msg)
+
+    taskiq_broker = BrokerWrapper(broker)
+
+    taskiq_broker.task(
+        "Hi!",
+        subject=subject,
+        schedule=[
+            {
+                "time": datetime.utcnow(),
+            },
+        ],
+    )
+
+    async with TestNatsBroker(broker):
+        task = asyncio.create_task(
+            run_scheduler(
+                SchedulerArgs(
+                    scheduler=TaskiqScheduler(
+                        broker=taskiq_broker,
+                        sources=[LabelScheduleSource(taskiq_broker)],
+                    ),
+                    modules=[],
+                ),
+            ),
+        )
+
+        with timeout_scope(3.0, True):
+            await event.wait()
+
+    mock.assert_called_once_with("Hi!")
+    task.cancel()

tests/test_rabbit.py

@@ -0,0 +1,63 @@
+import asyncio
+from datetime import datetime
+from unittest.mock import MagicMock
+
+import pytest
+from anyio import Event
+from faststream.rabbit import RabbitBroker, TestRabbitBroker
+from faststream.utils.functions import timeout_scope
+from taskiq import TaskiqScheduler
+from taskiq.cli.scheduler.args import SchedulerArgs
+from taskiq.cli.scheduler.run import run_scheduler
+from taskiq.schedule_sources import LabelScheduleSource
+
+from taskiq_faststream import BrokerWrapper
+
+
+@pytest.fixture
+def broker() -> RabbitBroker:
+    return RabbitBroker()
+
+
+@pytest.mark.anyio
+async def test_time_task(
+    subject: str,
+    broker: RabbitBroker,
+    mock: MagicMock,
+    event: Event,
+) -> None:
+    @broker.subscriber(subject)
+    async def handler(msg: str) -> None:
+        event.set()
+        mock(msg)
+
+    taskiq_broker = BrokerWrapper(broker)
+
+    taskiq_broker.task(
+        "Hi!",
+        queue=subject,
+        schedule=[
+            {
+                "time": datetime.utcnow(),
+            },
+        ],
+    )
+
+    async with TestRabbitBroker(broker):
+        task = asyncio.create_task(
+            run_scheduler(
+                SchedulerArgs(
+                    scheduler=TaskiqScheduler(
+                        broker=taskiq_broker,
+                        sources=[LabelScheduleSource(taskiq_broker)],
+                    ),
+                    modules=[],
+                ),
+            ),
+        )
+
+        with timeout_scope(3.0, True):
+            await event.wait()
+
+    mock.assert_called_once_with("Hi!")
+    task.cancel()