feat: add AppWrapper

Lancetnik · Lancetnik · commit a1c3ac2d4226 · 2023-11-12T16:07:57.000+03:00
diff --git a/taskiq_faststream/__init__.py b/taskiq_faststream/__init__.py
@@ -1,3 +1,6 @@
-from taskiq_faststream.broker import BrokerWrapper
+from taskiq_faststream.broker import AppWrapper, BrokerWrapper
 
-__all__ = ("BrokerWrapper",)
+__all__ = (
+    "BrokerWrapper",
+    "AppWrapper",
+)
diff --git a/taskiq_faststream/broker.py b/taskiq_faststream/broker.py
@@ -2,6 +2,7 @@
 import warnings
 
 from faststream._compat import TypeAlias, override
+from faststream.app import FastStream
 from faststream.broker.core.asyncronous import BrokerAsyncUsecase
 from faststream.types import SendableMessage
 from taskiq import AsyncBroker, BrokerMessage
@@ -44,10 +45,7 @@ async def shutdown(self) -> None:
 
     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)
+        await _broker_publish(self.broker, message)
 
     async def listen(
         self,
@@ -90,3 +88,49 @@ def task(  # type: ignore[override]
             schedule=schedule,
             **kwargs,
         )(lambda: None)
+
+
+class AppWrapper(BrokerWrapper):
+    """Wrap FastStream instance to taskiq compatible object.
+
+    Attributes:
+        app : FastStream instance.
+
+    Methods:
+        __init__ : Initializes the object.
+        startup : Startup wrapper FastStream.
+        shutdown : Shutdown wrapper FastStream.
+        kick : Call wrapped FastStream broker `publish` method.
+        task : Register FastStream scheduled task.
+    """
+
+    def __init__(self, app: FastStream) -> None:
+        super(BrokerWrapper, self).__init__()
+        self.app = app
+
+    async def startup(self) -> None:
+        """Startup wrapper FastStream broker."""
+        await super(BrokerWrapper, self).startup()
+        await self.app._startup()  # noqa: SLF001
+
+    async def shutdown(self) -> None:
+        """Shutdown wrapper FastStream broker."""
+        await self.app._shutdown()  # noqa: SLF001
+        await super(BrokerWrapper, self).shutdown()
+
+    async def kick(self, message: BrokerMessage) -> None:
+        """Call wrapped FastStream broker `publish` method."""
+        assert (  # noqa: S101
+            self.app.broker
+        ), "You should setup application broker firts"
+        await _broker_publish(self.app.broker, message)
+
+
+async def _broker_publish(
+    broker: BrokerAsyncUsecase[typing.Any, typing.Any],
+    message: BrokerMessage,
+) -> None:
+    labels = message.labels
+    labels.pop("schedule", None)
+    msg = await resolve_msg(labels.pop("message", message.message))
+    await broker.publish(msg, **labels)
diff --git a/tests/test_broker.py b/tests/test_broker.py
@@ -8,5 +8,5 @@ async def test_warning() -> None:
     broker = BrokerWrapper(None)
 
     with pytest.warns(RuntimeWarning):
-        async for _ in broker.listen():
+        async for _ in broker.listen():  # pragma: no branch
             break
diff --git a/tests/test_kafka.py b/tests/test_kafka.py
@@ -1,63 +1,25 @@
-import asyncio
-from datetime import datetime
-from unittest.mock import MagicMock
-
 import pytest
-from anyio import Event
+from faststream import FastStream
 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 import AsyncBroker
+
+from taskiq_faststream import AppWrapper
 
-from taskiq_faststream import BrokerWrapper
+from .testcase import SchedulerTestcase
 
 
 @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()
+class TestBroker(SchedulerTestcase):
+    test_class = TestKafkaBroker
+    subj_name = "topic"
+
+
+class TestApp(TestBroker):
+    @staticmethod
+    def build_taskiq_broker(broker: KafkaBroker) -> AsyncBroker:
+        """Build AppWrapper."""
+        return AppWrapper(FastStream(broker))
diff --git a/tests/test_nats.py b/tests/test_nats.py
@@ -1,63 +1,25 @@
-import asyncio
-from datetime import datetime
-from unittest.mock import MagicMock
-
 import pytest
-from anyio import Event
+from faststream import FastStream
 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 import AsyncBroker
+
+from taskiq_faststream import AppWrapper
 
-from taskiq_faststream import BrokerWrapper
+from .testcase import SchedulerTestcase
 
 
 @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()
+class TestBroker(SchedulerTestcase):
+    test_class = TestNatsBroker
+    subj_name = "subject"
+
+
+class TestApp(TestBroker):
+    @staticmethod
+    def build_taskiq_broker(broker: NatsBroker) -> AsyncBroker:
+        """Build AppWrapper."""
+        return AppWrapper(FastStream(broker))
diff --git a/tests/test_rabbit.py b/tests/test_rabbit.py
@@ -1,63 +1,25 @@
-import asyncio
-from datetime import datetime
-from unittest.mock import MagicMock
-
 import pytest
-from anyio import Event
+from faststream import FastStream
 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 import AsyncBroker
+
+from taskiq_faststream import AppWrapper
 
-from taskiq_faststream import BrokerWrapper
+from .testcase import SchedulerTestcase
 
 
 @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()
+class TestBroker(SchedulerTestcase):
+    test_class = TestRabbitBroker
+    subj_name = "queue"
+
+
+class TestApp(TestBroker):
+    @staticmethod
+    def build_taskiq_broker(broker: RabbitBroker) -> AsyncBroker:
+        """Build AppWrapper."""
+        return AppWrapper(FastStream(broker))
diff --git a/tests/testcase.py b/tests/testcase.py
@@ -0,0 +1,71 @@
+import asyncio
+from datetime import datetime
+from typing import Any
+from unittest.mock import MagicMock
+
+import pytest
+from anyio import Event
+from faststream.broker.core.asyncronous import BrokerAsyncUsecase
+from faststream.utils.functions import timeout_scope
+from taskiq import AsyncBroker, 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.mark.anyio
+class SchedulerTestcase:
+    test_class: Any
+    subj_name: str
+
+    @staticmethod
+    def build_taskiq_broker(broker: BrokerAsyncUsecase[Any, Any]) -> AsyncBroker:
+        """Build Taskiq compatible object."""
+        return BrokerWrapper(broker)
+
+    async def test_task(
+        self,
+        subject: str,
+        broker: BrokerAsyncUsecase[Any, Any],
+        mock: MagicMock,
+        event: Event,
+    ) -> None:
+        """Base testcase."""
+
+        @broker.subscriber(subject)
+        async def handler(msg: str) -> None:
+            event.set()
+            mock(msg)
+
+        taskiq_broker = self.build_taskiq_broker(broker)
+
+        taskiq_broker.task(
+            "Hi!",
+            **{self.subj_name: subject},
+            schedule=[
+                {
+                    "time": datetime.utcnow(),
+                },
+            ],
+        )
+
+        async with self.test_class(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()
