Repository: airtai/faststream (python)

Issue:
Feature: Allow broker creation in on_startup hook
For some reason, I need to delay the creation of an app's broker in an `on_startup` hook. This is mainly because my application has a pretty heavy warmup bootstrap that loads a lot of data into memory so I need to delay this bootstrap (which actually includes broker creation) after the parent process forked into worker subprocesses, otherwise the parent process has an unnecessary high memory footprint (that only worker processes need to have).

E.g.

```python
from faststream import FastStream
from faststream.kafka import KafkaBroker

app = FastStream()

@app.on_startup
def bootstrap():
    app.set_broker(KafkaBroker(["localhost:9092"]))
    # some warmup of the application pre-loading data into memory
    ...
```

For now, this isn't working because of this assert https://github.com/airtai/faststream/blob/9b7c33e8765485cf28f850cb485afcd838c12079/faststream/app.py#L40 which occurs at the very begining of the workers run, before `on_startup` hooks are executed.

The example above is overly simplified and one could argue that broker could have been created outside the `bootstrap()` function, but, as per example at https://faststream.airt.ai/latest/getting-started/lifespan/hooks/#lifespan, it could be legitimate to create it based on environment setting specifying not only the broker host but also the kind of broker to use (nats, kafka, ...).

I currently solve the problem by providing an *empty* `KafkaBroker` at app instanciation (`app = FastStream(broker=KafkaBroker())`) before overriding it using `app.set_broker()` in the `bootstrap()` function.

However, as per the base `Application` constructor at https://github.com/airtai/faststream/blob/9b7c33e8765485cf28f850cb485afcd838c12079/faststream/_internal/application.py#L48-L70 which allows application instanciation without a defined broker and this comment in `Application.set_broker()` https://github.com/airtai/faststream/blob/9b7c33e8765485cf28f850cb485afcd838c12079/faststream/_internal/application.py#L117-L120 which claims to be used for creation/init in `on_startup`, I would have expected my use case to be legitimate.

As a consequence, should this assert be removed, moved after startup hooks execution (at https://github.com/airtai/faststream/blob/9b7c33e8765485cf28f850cb485afcd838c12079/faststream/_internal/application.py#L168-L175) or just replaced by a warning when broker is `None` at https://github.com/airtai/faststream/blob/9b7c33e8765485cf28f850cb485afcd838c12079/faststream/_internal/application.py#L176-L177

Pull request:
feat: allow broker setting in on_startup hook
# Description

I've managed to achieve the expected behaviour mentioned in the issue #1884.

## Examples

Now it is possible to set broker in on_startup hook

```python
from faststream import FastStream
from faststream.kafka import KafkaBroker


app = FastStream()


@app.on_startup
def bootstrap():
    app.set_broker(KafkaBroker("localhost:9092"))
```

```sh
$ faststream run main:app
2025-02-18 18:26:41,955 INFO     - FastStream app starting...
2025-02-18 18:26:41,968 INFO     - FastStream app started successfully! To exit, press CTRL+C
```

If the broke was not set the exception will be raised

```python
from faststream import FastStream
from faststream.kafka import KafkaBroker


app = FastStream()
```

```sh
$ faststream run main:app
...
AssertionError: You should setup a broker
```

I had to refactor asgi tests a little bit. As this change requires to pass broker to the application, I had to create test brokers to pass all tests.

Fixes #1884

## Type of change

Please delete options that are not relevant.

- [x] New feature (a non-breaking change that adds functionality)

## Checklist

- [x] My code adheres to the style guidelines of this project (`scripts/lint.sh` shows no errors)
- [x] I have conducted a self-review of my own code
- [x] I have made the necessary changes to the documentation
- [x] My changes do not generate any new warnings
- [x] I have added tests to validate the effectiveness of my fix or the functionality of my new feature
- [x] Both new and existing unit tests pass successfully on my local environment by running `scripts/test-cov.sh`
- [x] I have ensured that static analysis tests are passing by running `scripts/static-analysis.sh`
- [x] I have included code examples to illustrate the modifications


Patch:
diff --git a/faststream/_internal/application.py b/faststream/_internal/application.py
index ff2e39fa..1074c44c 100644
--- a/faststream/_internal/application.py
+++ b/faststream/_internal/application.py
@@ -197,6 +197,7 @@ class Application(ABC, AsyncAPIApplication):
     ) -> None:
         self._log(log_level, "FastStream app starting...")
         await self.start(**(run_extra_options or {}))
+        assert self.broker, "You should setup a broker"  # nosec B101
         self._log(
             log_level, "FastStream app started successfully! To exit, press CTRL+C"
         )
diff --git a/faststream/app.py b/faststream/app.py
index fc18797f..67d9be76 100644
--- a/faststream/app.py
+++ b/faststream/app.py
@@ -37,8 +37,6 @@ class FastStream(Application):
         sleep_time: float = 0.1,
     ) -> None:
         """Run FastStream Application."""
-        assert self.broker, "You should setup a broker"  # nosec B101
-
         set_exit(lambda *_: self.exit(), sync=False)
 
         async with catch_startup_validation_error(), self.lifespan_context(


Test patch:
diff --git a/tests/asgi/testcase.py b/tests/asgi/testcase.py
index fccfedb7..0407e9d7 100644
--- a/tests/asgi/testcase.py
+++ b/tests/asgi/testcase.py
@@ -1,4 +1,5 @@
 from typing import Any
+from unittest.mock import AsyncMock
 
 import pytest
 from starlette.testclient import TestClient
@@ -14,33 +15,45 @@ class AsgiTestcase:
     def get_test_broker(self, broker) -> Any:
         raise NotImplementedError()
 
-    def test_not_found(self):
-        app = AsgiFastStream()
+    @pytest.mark.asyncio
+    async def test_not_found(self):
+        broker = self.get_broker()
+        app = AsgiFastStream(broker)
 
-        with TestClient(app) as client:
-            response = client.get("/")
-            assert response.status_code == 404
+        async with self.get_test_broker(broker):
+            with TestClient(app) as client:
+                response = client.get("/")
+                assert response.status_code == 404
 
-    def test_ws_not_found(self):
-        app = AsgiFastStream()
+    @pytest.mark.asyncio
+    async def test_ws_not_found(self):
+        broker = self.get_broker()
+
+        app = AsgiFastStream(broker)
 
-        with TestClient(app) as client:  # noqa: SIM117
-            with pytest.raises(WebSocketDisconnect):
-                with client.websocket_connect("/ws"):  # raises error
-                    pass
+        async with self.get_test_broker(broker):
+            with TestClient(app) as client:
+                with pytest.raises(WebSocketDisconnect):
+                    with client.websocket_connect("/ws"):  # raises error
+                        pass
 
-    def test_asgi_ping_unhealthy(self):
+    @pytest.mark.asyncio
+    async def test_asgi_ping_unhealthy(self):
         broker = self.get_broker()
 
         app = AsgiFastStream(
+            broker,
             asgi_routes=[
                 ("/health", make_ping_asgi(broker, timeout=5.0)),
-            ]
+            ],
         )
+        async with self.get_test_broker(broker) as br:
+            br.ping = AsyncMock()
+            br.ping.return_value = False
 
-        with TestClient(app) as client:
-            response = client.get("/health")
-            assert response.status_code == 500
+            with TestClient(app) as client:
+                response = client.get("/health")
+                assert response.status_code == 500
 
     @pytest.mark.asyncio
     async def test_asgi_ping_healthy(self):
@@ -68,14 +81,17 @@ class AsgiTestcase:
                 assert response.status_code == 200
                 assert response.text
 
-    def test_get_decorator(self):
+    @pytest.mark.asyncio
+    async def test_get_decorator(self):
         @get
         async def some_handler(scope):
             return AsgiResponse(body=b"test", status_code=200)
 
-        app = AsgiFastStream(asgi_routes=[("/test", some_handler)])
+        broker = self.get_broker()
+        app = AsgiFastStream(broker, asgi_routes=[("/test", some_handler)])
 
-        with TestClient(app) as client:
-            response = client.get("/test")
-            assert response.status_code == 200
-            assert response.text == "test"
+        async with self.get_test_broker(broker):
+            with TestClient(app) as client:
+                response = client.get("/test")
+                assert response.status_code == 200
+                assert response.text == "test"
diff --git a/tests/cli/rabbit/test_app.py b/tests/cli/rabbit/test_app.py
index 68fc2517..5db19a9e 100644
--- a/tests/cli/rabbit/test_app.py
+++ b/tests/cli/rabbit/test_app.py
@@ -10,6 +10,7 @@ import pytest
 from faststream import FastStream, TestApp
 from faststream._compat import IS_WINDOWS
 from faststream.log import logger
+from faststream.rabbit.testing import TestRabbitBroker
 
 
 def test_init(app: FastStream, broker):
@@ -31,6 +32,23 @@ def test_set_broker(broker, app_without_broker: FastStream):
     assert app_without_broker.broker is broker
 
 
+@pytest.mark.asyncio
+async def test_set_broker_in_on_startup_hook(app_without_broker: FastStream, broker):
+    def add_broker():
+        app_without_broker.set_broker(broker)
+
+    app_without_broker.on_startup(add_broker)
+
+    async with TestRabbitBroker(broker):
+        await app_without_broker._startup()
+
+
+@pytest.mark.asyncio
+async def test_startup_fails_if_no_broker_was_provided(app_without_broker: FastStream):
+    with pytest.raises(AssertionError):
+        await app_without_broker._startup()
+
+
 def test_log(app: FastStream, app_without_logger: FastStream):
     app._log(logging.INFO, "test")
     app_without_logger._log(logging.INFO, "test")