Feature: Implementation of the AsyncContextManager interface for Application #3228

docs/docs/en/getting-started/lifespan/lifecycle.md

@@ -23,6 +23,9 @@ Most of the time you start an application with `faststream run module:app` and n
 | `stop()`  | Runs the `on_shutdown` hooks, waits for in-flight handlers, stops every broker, runs `after_shutdown`. |
 | `exit()`  | Asks a running `run()` to stop.                                                               |
 
+!!! note
+    You can use `async with app:` instead of calling `start()` and `stop()` manually.
+
 ## Running Without the CLI
 
 `faststream run module:app` imports the module and awaits `#!python app.run()`, passing the CLI's log level and extra options along. You can do the same yourself (if your hooks take CLI options, pass them through `run_extra_options=`):

faststream/_internal/application.py

@@ -4,7 +4,7 @@
 from contextlib import asynccontextmanager
 from typing import TYPE_CHECKING, Any, Optional, TypeVar
 
-from typing_extensions import ParamSpec, deprecated
+from typing_extensions import ParamSpec, Self, deprecated
 
 from faststream._internal.di import FastDependsConfig
 from faststream._internal.logger import logger
@@ -14,6 +14,8 @@
 from faststream.specification import AsyncAPI
 
 if TYPE_CHECKING:
+    from types import TracebackType
+
     from faststream._internal.basic_types import (
         AnyCallable,
         AsyncFunc,
@@ -191,6 +193,19 @@ def __init__(
         else:
             self.lifespan_context = fake_context
 
+    async def __aenter__(self, **run_extra_options: "SettingField") -> Self:
+        await self.start(**run_extra_options)
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None = None,
+        exc_val: BaseException | None = None,
+        exc_tb: Optional["TracebackType"] = None,
+        **run_extra_options: "SettingField",
+    ) -> None:
+        await self.stop()
+
     @abstractmethod
     def exit(self) -> None:
         """Stop application manually."""

faststream/_internal/testing/app.py

@@ -33,11 +33,11 @@ def __enter__(self) -> "Application":
 
             lifespan_context = self.app.lifespan_context(**self._extra_options)
             stack.enter_context(portal.wrap_async_context_manager(lifespan_context))
-            portal.call(partial(self.app.start, **self._extra_options))
+            portal.call(partial(self.app.__aenter__, **self._extra_options))
 
             @stack.callback
             def wait_shutdown() -> None:
-                portal.call(self.app.stop)
+                portal.call(partial(self.app.__aexit__, **self._extra_options))
 
             self.exit_stack = stack.pop_all()
 
@@ -54,8 +54,7 @@ def __exit__(
     async def __aenter__(self) -> "Application":
         self.lifespan_scope = self.app.lifespan_context(**self._extra_options)
         await self.lifespan_scope.__aenter__()
-        await self.app.start(**self._extra_options)
-        return self.app
+        return await self.app.__aenter__(**self._extra_options)
 
     async def __aexit__(
         self,
@@ -64,5 +63,5 @@ async def __aexit__(
         exc_tb: Optional["TracebackType"] = None,
     ) -> None:
         """Exit the asynchronous context manager."""
-        await self.app.stop()
+        await self.app.__aexit__(exc_type, exc_val, exc_tb)
         await self.lifespan_scope.__aexit__(exc_type, exc_val, exc_tb)

tests/asgi/test_app.py

@@ -9,6 +9,7 @@
 
 from collections.abc import Callable
 from typing import Any
+from unittest.mock import MagicMock
 
 import pytest
 from fast_depends import Depends
@@ -27,6 +28,7 @@
 )
 from faststream.asgi.params import Header, Query
 from faststream.asgi.types import ASGIApp, Scope
+from faststream.rabbit import RabbitBroker
 from faststream.specification import AsyncAPI
 
 
@@ -99,6 +101,21 @@ def test_asyncapi_json_disabled() -> None:
     assert client.get("/docs.json").status_code == 404
 
 
+@pytest.mark.asyncio()
+async def test_async_context_manager(mock: MagicMock) -> None:
+    app = AsgiFastStream(
+        MagicMock(spec=RabbitBroker),
+        on_startup=[mock.on],
+        on_shutdown=[mock.off],
+    )
+
+    async with app as context_app:
+        assert context_app is app
+
+    mock.on.assert_called_once()
+    mock.off.assert_called_once()
+
+
 @pytest.mark.parametrize(
     ("decorator", "client_method"),
     (

tests/cli/require_broker/test_app.py

@@ -313,6 +313,35 @@ async def test_test_app(mock: MagicMock) -> None:
     mock.off.assert_called_once()
 
 
+@pytest.mark.asyncio()
+async def test_application_async_context_manager(mock: MagicMock) -> None:
+    app = FastStream(AsyncMock(spec=RabbitBroker))
+
+    app.on_startup(mock.on)
+    app.on_shutdown(mock.off)
+
+    async with app as context_app:
+        assert context_app is app
+
+    mock.on.assert_called_once()
+    mock.off.assert_called_once()
+
+
+@pytest.mark.asyncio()
+async def test_application_async_context_manager_stops_on_exception(
+    mock: MagicMock,
+) -> None:
+    app = FastStream(AsyncMock(spec=RabbitBroker))
+    app.on_shutdown(mock.off)
+    error = ValueError("expected")
+
+    with pytest.raises(ValueError, match="expected"):
+        async with app:
+            raise error
+
+    mock.off.assert_called_once()
+
+
 @pytest.mark.asyncio()
 async def test_test_app_with_excp(mock: MagicMock) -> None:
     app = FastStream(AsyncMock(spec=RabbitBroker))