feat: delayed broker setting (#1931)

* feat: delayed broker setting

* feat: tests for delayed broker setting

* feat: addition test

* feat: broker -> **brokers

* feat: ruff

faststream/_internal/application.py

@@ -29,6 +29,7 @@
     fake_context,
     to_async,
 )
+from faststream.exceptions import SetupError
 
 if TYPE_CHECKING:
     from fast_depends.library.serializer import SerializerProto
@@ -78,7 +79,7 @@ async def catch_startup_validation_error() -> AsyncIterator[None]:
 class StartAbleApplication:
     def __init__(
         self,
-        broker: "BrokerUsecase[Any, Any]",
+        broker: Optional["BrokerUsecase[Any, Any]"] = None,
         /,
         provider: Optional["Provider"] = None,
         serializer: Optional["SerializerProto"] = EMPTY,
@@ -91,7 +92,7 @@ def __init__(
 
     def _init_setupable_(  # noqa: PLW3201
         self,
-        broker: "BrokerUsecase[Any, Any]",
+        broker: Optional["BrokerUsecase[Any, Any]"] = None,
         /,
         provider: Optional["Provider"] = None,
         serializer: Optional["SerializerProto"] = EMPTY,
@@ -115,21 +116,39 @@ def _init_setupable_(  # noqa: PLW3201
             )
         )
 
-        self.broker = broker
+        self.brokers = [broker] if broker else []
 
         self._setup()
 
     def _setup(self) -> None:
-        self.broker._setup(OuterBrokerState(di_state=self._state.di_state))
+        for broker in self.brokers:
+            broker._setup(OuterBrokerState(di_state=self._state.di_state))
 
     async def _start_broker(self) -> None:
+        assert self.broker, "You should setup a broker"
         await self.broker.start()
 
+    @property
+    def broker(self) -> Optional["BrokerUsecase[Any, Any]"]:
+        return self.brokers[0] if self.brokers else None
+
+    def set_broker(self, broker: "BrokerUsecase[Any, Any]") -> None:
+        """Set already existed App object broker.
+
+        Useful then you create/init broker in `on_startup` hook.
+        """
+        if self.brokers:
+            msg = f"`{self}` already has a broker. You can't use multiple brokers until 1.0.0 release."
+            raise SetupError(msg)
+
+        self.brokers.append(broker)
+        self._setup()
+
 
 class Application(StartAbleApplication):
     def __init__(
         self,
-        broker: "BrokerUsecase[Any, Any]",
+        broker: Optional["BrokerUsecase[Any, Any]"] = None,
         /,
         logger: Optional["LoggerProto"] = logger,
         provider: Optional["Provider"] = None,
@@ -253,7 +272,8 @@ async def _shutdown(self, log_level: int = logging.INFO) -> None:
     async def stop(self) -> None:
         """Executes shutdown hooks and stop broker."""
         async with self._shutdown_hooks_context():
-            await self.broker.close()
+            for broker in self.brokers:
+                await broker.close()
 
     @asynccontextmanager
     async def _shutdown_hooks_context(self) -> AsyncIterator[None]:

faststream/_internal/cli/utils/logs.py

@@ -68,4 +68,5 @@ def set_log_level(level: int, app: "Application") -> None:
     if app.logger and getattr(app.logger, "setLevel", None):
         app.logger.setLevel(level)  # type: ignore[attr-defined]
 
-    app.broker._state.get().logger_state.set_level(level)
+    for broker in app.brokers:
+        broker._state.get().logger_state.set_level(level)

faststream/asgi/app.py

@@ -82,7 +82,7 @@ class AsgiFastStream(Application):
 
     def __init__(
         self,
-        broker: "BrokerUsecase[Any, Any]",
+        broker: Optional["BrokerUsecase[Any, Any]"] = None,
         /,
         asgi_routes: Sequence[tuple[str, "ASGIApp"]] = (),
         # regular broker args

tests/application/test_delayed_broker.py

@@ -0,0 +1,38 @@
+import pytest
+
+from faststream._internal.application import StartAbleApplication
+from faststream.exceptions import SetupError
+from faststream.rabbit import RabbitBroker
+
+
+def test_set_broker() -> None:
+    app = StartAbleApplication()
+
+    assert app.broker is None
+
+    broker = RabbitBroker()
+    app.set_broker(broker)
+
+    assert app.broker is broker
+
+
+def test_set_more_than_once_broker() -> None:
+    app = StartAbleApplication()
+    broker_1 = RabbitBroker()
+    broker_2 = RabbitBroker()
+
+    app.set_broker(broker_1)
+
+    with pytest.raises(
+        SetupError,
+        match=f"`{app}` already has a broker. You can't use multiple brokers until 1.0.0 release.",
+    ):
+        app.set_broker(broker_2)
+
+
+@pytest.mark.asyncio()
+async def test_start_not_setup_broker() -> None:
+    app = StartAbleApplication()
+
+    with pytest.raises(AssertionError, match="You should setup a broker"):
+        await app._start_broker()