Add support for lifespan state

mangum/adapter.py

@@ -72,12 +72,14 @@ def infer(self, event: LambdaEvent, context: LambdaContext) -> LambdaHandler:
 
     def __call__(self, event: LambdaEvent, context: LambdaContext) -> dict:
         handler = self.infer(event, context)
+        scope = handler.scope
         with ExitStack() as stack:
             if self.lifespan in ("auto", "on"):
                 lifespan_cycle = LifespanCycle(self.app, self.lifespan)
                 stack.enter_context(lifespan_cycle)
+                scope.update({"state": lifespan_cycle.lifespan_state.copy()})
 
-            http_cycle = HTTPCycle(handler.scope, handler.body)
+            http_cycle = HTTPCycle(scope, handler.body)
             http_response = http_cycle(self.app)
 
             return handler(http_response)

mangum/protocols/lifespan.py

@@ -2,7 +2,7 @@
 import enum
 import logging
 from types import TracebackType
-from typing import Optional, Type
+from typing import Any, Dict, Optional, Type
 
 from mangum.types import ASGI, LifespanMode, Message
 from mangum.exceptions import LifespanUnsupported, LifespanFailure, UnexpectedMessage
@@ -62,6 +62,7 @@ def __init__(self, app: ASGI, lifespan: LifespanMode) -> None:
         self.startup_event: asyncio.Event = asyncio.Event()
         self.shutdown_event: asyncio.Event = asyncio.Event()
         self.logger = logging.getLogger("mangum.lifespan")
+        self.lifespan_state: Dict[str, Any] = {}
 
     def __enter__(self) -> None:
         """Runs the event loop for application startup."""
@@ -81,7 +82,7 @@ async def run(self) -> None:
         """Calls the application with the `lifespan` connection scope."""
         try:
             await self.app(
-                {"type": "lifespan", "asgi": {"spec_version": "2.0", "version": "3.0"}},
+                {"type": "lifespan", "asgi": {"spec_version": "2.0", "version": "3.0"}, "state": self.lifespan_state},
                 self.receive,
                 self.send,
             )

tests/test_lifespan.py

@@ -211,6 +211,58 @@ async def app(scope, receive, send):
         handler(mock_aws_api_gateway_event, {})
 
 
+@pytest.mark.parametrize(
+    "mock_aws_api_gateway_event,lifespan_state,lifespan",
+    [
+        (["GET", None, None], {"test_key": "test_value"}, "auto"),
+        (["GET", None, None], {"test_key": "test_value"}, "on"),
+    ],
+    indirect=["mock_aws_api_gateway_event"],
+)
+def test_lifespan_state(mock_aws_api_gateway_event, lifespan_state, lifespan) -> None:
+    startup_complete = False
+    shutdown_complete = False
+
+    async def app(scope, receive, send):
+        nonlocal startup_complete, shutdown_complete
+
+        if scope["type"] == "lifespan":
+            while True:
+                message = await receive()
+                if message["type"] == "lifespan.startup":
+                    scope["state"].update(lifespan_state)
+                    await send({"type": "lifespan.startup.complete"})
+                    startup_complete = True
+                elif message["type"] == "lifespan.shutdown":
+                    await send({"type": "lifespan.shutdown.complete"})
+                    shutdown_complete = True
+                    return
+
+        if scope["type"] == "http":
+            assert lifespan_state.items() <= scope["state"].items()
+            await send(
+                {
+                    "type": "http.response.start",
+                    "status": 200,
+                    "headers": [[b"content-type", b"text/plain; charset=utf-8"]],
+                }
+            )
+            await send({"type": "http.response.body", "body": b"Hello, world!"})
+
+    handler = Mangum(app, lifespan=lifespan)
+    response = handler(mock_aws_api_gateway_event, {})
+
+    assert startup_complete
+    assert shutdown_complete
+    assert response == {
+        "statusCode": 200,
+        "isBase64Encoded": False,
+        "headers": {"content-type": "text/plain; charset=utf-8"},
+        "multiValueHeaders": {},
+        "body": "Hello, world!",
+    }
+
+
 @pytest.mark.parametrize("mock_aws_api_gateway_event", [["GET", None, None]], indirect=True)
 def test_starlette_lifespan(mock_aws_api_gateway_event) -> None:
     startup_complete = False