fix: accept an app whose broker is created in an on_startup hook

README.md

@@ -73,7 +73,7 @@ The current `StreamMessage` is resolvable within DI via the pre-built `faststrea
 
 ## API
 
-- `setup_di(app, container)` — stores the container in the app context and registers startup/shutdown lifecycle hooks: on startup it reopens the container and adds the DI middleware to every broker of the app (including one added via `app.add_broker` after `setup_di`), after shutdown it closes the container
+- `setup_di(app, container)` — stores the container in the app context and registers startup/shutdown lifecycle hooks: on startup it reopens the container and adds the DI middleware to every broker of the app (including one added via `app.add_broker` after `setup_di`, or by an `on_startup` hook registered before `setup_di`), after shutdown it closes the container. Raises at startup if the app has no broker by then
 - `FromDI(dependency, *, use_cache=True, cast=False)` — FastStream `Depends` that resolves a provider (or type) from the request container. Raises `RuntimeError` naming `setup_di` when a message reaches it without the middleware installed
 - `fetch_di_container(app)` — returns the root container from the app context
 - `faststream_message_provider` — `ContextProvider` for the current `faststream.StreamMessage`

docs/adr/0002-install-middleware-on-startup.md

@@ -30,9 +30,14 @@ from, rather than a private record of installed brokers that could drift from it
 Between `setup_di` and startup the middleware is not yet on any broker. Nothing in this package,
 its tests, or its documentation inspects a broker in that window; `TestApp` runs the startup hooks.
 
-The `if not app.broker` guard in `setup_di` is kept as it was. Relaxing it so that a broker
-created inside the user's own `on_startup` hook can be picked up is possible, since hooks run in
-registration order, but that ordering is subtle enough to want its own decision.
+`setup_di` no longer requires a broker at call time. The original version of this decision kept
+the `if not app.broker` guard; [#56](https://github.com/modern-python/modern-di-faststream/issues/56)
+dropped it so that a broker created inside the user's own `on_startup` hook is covered, because the
+broker list is read at startup anyway. Hooks run in registration order, so that hook must be
+registered before `setup_di`; the install hook raises when the list is still empty when it runs,
+naming both remedies, and the message-time error from `FromDI` names the other order. A broker
+that a later hook adds while another broker already exists is the one case that still surfaces at
+message time.
 
 **Revisit trigger:** FastStream exposes a hook for a broker being added to an app, so the
 middleware can be installed at that moment instead of on startup, **or** FastStream freezes a

modern_di_faststream/main.py

@@ -22,7 +22,13 @@
 _MISSING_REQUEST_CONTAINER = (
     "No request container for this message, so the DI middleware did not run for it. "
     "Call setup_di(app, container) on the app that owns this broker and start the app; "
-    "in tests, pair the test broker with TestApp(app) in the same `async with`."
+    "in tests, pair the test broker with TestApp(app) in the same `async with`. "
+    "A broker added by an on_startup hook is covered only if that hook was registered before setup_di."
+)
+_NO_BROKER_AT_STARTUP = (
+    "No broker on the app when the DI middleware is installed at startup. "
+    "Pass one to FastStream(...) or app.add_broker(...) before startup; "
+    "an on_startup hook that adds it must be registered before setup_di, since hooks run in registration order."
 )
 
 
@@ -65,17 +71,15 @@ def setup_di(
     app: faststream.FastStream | AsgiFastStream,
     container: Container,
 ) -> Container:
-    if not app.broker:
-        msg = "Broker must be defined to setup DI"
-        raise RuntimeError(msg)
-
     container.add_providers(faststream_message_provider)
     app.context.set_global(_ROOT_CONTAINER_KEY, container)
     middleware_factory = _DIMiddlewareFactory(container)
 
     # Installed on startup rather than here so a broker added after ``setup_di`` is covered,
     # and skipped where present so a restart adds no second copy: docs/adr/0002-install-middleware-on-startup.md
     def install_middleware() -> None:
+        if not app.brokers:
+            raise RuntimeError(_NO_BROKER_AT_STARTUP)
         for broker in app.brokers:
             if middleware_factory not in broker.config.broker_middlewares:
                 broker.add_middleware(middleware_factory)

tests/test_faststream_di.py

@@ -48,9 +48,38 @@ async def index_subscriber(
         assert result_str == b""
 
 
-async def test_app_without_broker() -> None:
-    with pytest.raises(RuntimeError, match="Broker must be defined to setup DI"):
-        modern_di_faststream.setup_di(faststream.FastStream(), container=Container())
+async def test_broker_added_by_a_startup_hook_registered_before_setup_di_gets_di() -> None:
+    """INVARIANT: ``setup_di`` accepts an app whose broker is created in an ``on_startup`` hook.
+
+    Broken by refusing an app with no broker at ``setup_di`` time, or by installing the middleware
+    anywhere but a startup hook registered by ``setup_di`` itself: hooks run in registration order,
+    so the user's ``add_broker`` hook has run by the time the install hook reads ``app.brokers``.
+    FastStream documents this shape (``set_broker``: "create/init broker in ``on_startup`` hook").
+    """
+    broker = NatsBroker()
+    app_ = faststream.FastStream()
+    resolved: list[SimpleCreator] = []
+
+    @app_.on_startup
+    async def attach_broker() -> None:
+        app_.add_broker(broker)
+
+    modern_di_faststream.setup_di(app_, container=Container(groups=[Dependencies]))
+    _subscribe_resolving(broker, TEST_SUBJECT, resolved)
+
+    async with TestNatsBroker(broker) as br, TestApp(app_):
+        await br.publish(None, TEST_SUBJECT)
+
+    assert [type(instance) for instance in resolved] == [SimpleCreator]
+
+
+async def test_app_without_broker_at_startup_names_both_remedies() -> None:
+    """The empty broker list is reported when the install hook runs, before FastStream's own assert."""
+    app_ = faststream.FastStream()
+    modern_di_faststream.setup_di(app_, container=Container())
+
+    with pytest.raises(RuntimeError, match=r"add_broker.*before setup_di"):
+        await app_.start()
 
 
 def test_fetch_di_container(app: faststream.FastStream) -> None: