From 1e4ace1797ec5339bc890367962f75bb8b7f63e4 Mon Sep 17 00:00:00 2001 From: Muhammad Furqan Habibi Date: Fri, 7 Jun 2024 17:49:06 +0900 Subject: [PATCH 1/8] Add support for lifespan state --- mangum/protocols/lifespan.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/mangum/protocols/lifespan.py b/mangum/protocols/lifespan.py index 22ed016..1e495c8 100644 --- a/mangum/protocols/lifespan.py +++ b/mangum/protocols/lifespan.py @@ -4,6 +4,7 @@ import enum import logging from types import TracebackType +from typing import Any from mangum.exceptions import LifespanFailure, LifespanUnsupported, UnexpectedMessage from mangum.types import ASGI, LifespanMode, Message @@ -63,6 +64,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.""" @@ -82,7 +84,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, ) @@ -101,7 +103,7 @@ async def receive(self) -> Message: if self.state is LifespanCycleState.CONNECTING: # Connection established. The next event returned by the queue will be # `lifespan.startup` to inform the application that the connection is - # ready to receive lfiespan messages. + # ready to receive lifespan messages. self.state = LifespanCycleState.STARTUP elif self.state is LifespanCycleState.STARTUP: From 3c3f204e899bb54e61fcbaa66caace9f136dd421 Mon Sep 17 00:00:00 2001 From: Muhammad Furqan Habibi Date: Tue, 11 Jun 2024 14:08:29 +0900 Subject: [PATCH 2/8] Add lifespan state to http connection --- mangum/adapter.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/mangum/adapter.py b/mangum/adapter.py index 0bbb10c..e21bbe3 100644 --- a/mangum/adapter.py +++ b/mangum/adapter.py @@ -60,12 +60,14 @@ def infer(self, event: LambdaEvent, context: LambdaContext) -> LambdaHandler: def __call__(self, event: LambdaEvent, context: LambdaContext) -> dict[str, Any]: 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 |= {"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) From 56a9a58a6232feb8951ee97a7c0f8375e1dcfdad Mon Sep 17 00:00:00 2001 From: Muhammad Furqan Habibi Date: Thu, 13 Jun 2024 12:05:13 +0900 Subject: [PATCH 3/8] Add test --- mangum/adapter.py | 2 +- mangum/protocols/http.py | 10 ++++++-- tests/test_lifespan.py | 52 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 61 insertions(+), 3 deletions(-) diff --git a/mangum/adapter.py b/mangum/adapter.py index e21bbe3..4cbfeba 100644 --- a/mangum/adapter.py +++ b/mangum/adapter.py @@ -65,7 +65,7 @@ def __call__(self, event: LambdaEvent, context: LambdaContext) -> dict[str, Any] if self.lifespan in ("auto", "on"): lifespan_cycle = LifespanCycle(self.app, self.lifespan) stack.enter_context(lifespan_cycle) - scope |= {"state": lifespan_cycle.lifespan_state.copy()} + scope.update({"state": lifespan_cycle.lifespan_state.copy()}) http_cycle = HTTPCycle(scope, handler.body) http_response = http_cycle(self.app) diff --git a/mangum/protocols/http.py b/mangum/protocols/http.py index f0b50e5..6ab6b37 100644 --- a/mangum/protocols/http.py +++ b/mangum/protocols/http.py @@ -82,11 +82,17 @@ async def receive(self) -> Message: return await self.app_queue.get() # pragma: no cover async def send(self, message: Message) -> None: - if self.state is HTTPCycleState.REQUEST and message["type"] == "http.response.start": + if ( + self.state is HTTPCycleState.REQUEST + and message["type"] == "http.response.start" + ): self.status = message["status"] self.headers = message.get("headers", []) self.state = HTTPCycleState.RESPONSE - elif self.state is HTTPCycleState.RESPONSE and message["type"] == "http.response.body": + elif ( + self.state is HTTPCycleState.RESPONSE + and message["type"] == "http.response.body" + ): body = message.get("body", b"") more_body = message.get("more_body", False) self.buffer.write(body) diff --git a/tests/test_lifespan.py b/tests/test_lifespan.py index 42cb024..12a98ef 100644 --- a/tests/test_lifespan.py +++ b/tests/test_lifespan.py @@ -209,6 +209,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 From 737dc60f6630c2ebd2c72b39bb05bbc451e16b25 Mon Sep 17 00:00:00 2001 From: Muhammad Furqan Habibi Date: Thu, 26 Sep 2024 13:26:38 +0900 Subject: [PATCH 4/8] Revert import changes --- mangum/protocols/http.py | 12 +++--------- mangum/protocols/lifespan.py | 2 +- tests/test_lifespan.py | 4 +++- 3 files changed, 7 insertions(+), 11 deletions(-) diff --git a/mangum/protocols/http.py b/mangum/protocols/http.py index 6ab6b37..fc95787 100644 --- a/mangum/protocols/http.py +++ b/mangum/protocols/http.py @@ -3,8 +3,8 @@ import logging from io import BytesIO +from mangum.types import ASGI, Message, Scope, Response from mangum.exceptions import UnexpectedMessage -from mangum.types import ASGI, Message, Response, Scope class HTTPCycleState(enum.Enum): @@ -82,17 +82,11 @@ async def receive(self) -> Message: return await self.app_queue.get() # pragma: no cover async def send(self, message: Message) -> None: - if ( - self.state is HTTPCycleState.REQUEST - and message["type"] == "http.response.start" - ): + if self.state is HTTPCycleState.REQUEST and message["type"] == "http.response.start": self.status = message["status"] self.headers = message.get("headers", []) self.state = HTTPCycleState.RESPONSE - elif ( - self.state is HTTPCycleState.RESPONSE - and message["type"] == "http.response.body" - ): + elif self.state is HTTPCycleState.RESPONSE and message["type"] == "http.response.body": body = message.get("body", b"") more_body = message.get("more_body", False) self.buffer.write(body) diff --git a/mangum/protocols/lifespan.py b/mangum/protocols/lifespan.py index 1e495c8..923a660 100644 --- a/mangum/protocols/lifespan.py +++ b/mangum/protocols/lifespan.py @@ -6,8 +6,8 @@ from types import TracebackType from typing import Any -from mangum.exceptions import LifespanFailure, LifespanUnsupported, UnexpectedMessage from mangum.types import ASGI, LifespanMode, Message +from mangum.exceptions import LifespanUnsupported, LifespanFailure, UnexpectedMessage class LifespanCycleState(enum.Enum): diff --git a/tests/test_lifespan.py b/tests/test_lifespan.py index 12a98ef..a769bac 100644 --- a/tests/test_lifespan.py +++ b/tests/test_lifespan.py @@ -1,13 +1,15 @@ import logging import pytest -from quart import Quart + from starlette.applications import Starlette from starlette.responses import PlainTextResponse from mangum import Mangum from mangum.exceptions import LifespanFailure +from quart import Quart + @pytest.mark.parametrize( "mock_aws_api_gateway_event,lifespan", From 5f0111a88df39e47726b7a32069b77086f6e2e52 Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Thu, 26 Sep 2024 22:21:33 +0200 Subject: [PATCH 5/8] Fix imports --- mangum/protocols/lifespan.py | 2 +- tests/test_lifespan.py | 4 +--- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/mangum/protocols/lifespan.py b/mangum/protocols/lifespan.py index 923a660..1e495c8 100644 --- a/mangum/protocols/lifespan.py +++ b/mangum/protocols/lifespan.py @@ -6,8 +6,8 @@ from types import TracebackType from typing import Any +from mangum.exceptions import LifespanFailure, LifespanUnsupported, UnexpectedMessage from mangum.types import ASGI, LifespanMode, Message -from mangum.exceptions import LifespanUnsupported, LifespanFailure, UnexpectedMessage class LifespanCycleState(enum.Enum): diff --git a/tests/test_lifespan.py b/tests/test_lifespan.py index a769bac..12a98ef 100644 --- a/tests/test_lifespan.py +++ b/tests/test_lifespan.py @@ -1,15 +1,13 @@ import logging import pytest - +from quart import Quart from starlette.applications import Starlette from starlette.responses import PlainTextResponse from mangum import Mangum from mangum.exceptions import LifespanFailure -from quart import Quart - @pytest.mark.parametrize( "mock_aws_api_gateway_event,lifespan", From 01a96a83957ca8e6317c6e0e5454428853ec8f77 Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Thu, 26 Sep 2024 22:21:59 +0200 Subject: [PATCH 6/8] Fix imports --- mangum/protocols/http.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mangum/protocols/http.py b/mangum/protocols/http.py index fc95787..f0b50e5 100644 --- a/mangum/protocols/http.py +++ b/mangum/protocols/http.py @@ -3,8 +3,8 @@ import logging from io import BytesIO -from mangum.types import ASGI, Message, Scope, Response from mangum.exceptions import UnexpectedMessage +from mangum.types import ASGI, Message, Response, Scope class HTTPCycleState(enum.Enum): From 8188ab8b82dc72c9e5c1b2e55f26cf0a79c92807 Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Thu, 26 Sep 2024 22:26:22 +0200 Subject: [PATCH 7/8] Improve a bit the test --- tests/test_lifespan.py | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/tests/test_lifespan.py b/tests/test_lifespan.py index 12a98ef..ef1de7d 100644 --- a/tests/test_lifespan.py +++ b/tests/test_lifespan.py @@ -1,4 +1,5 @@ import logging +from typing import Literal import pytest from quart import Quart @@ -7,6 +8,7 @@ from mangum import Mangum from mangum.exceptions import LifespanFailure +from mangum.types import Receive, Scope, Send @pytest.mark.parametrize( @@ -210,25 +212,22 @@ async def app(scope, receive, send): @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"), - ], + "mock_aws_api_gateway_event,lifespan", + [(["GET", None, None], "auto"), (["GET", None, None], "on")], indirect=["mock_aws_api_gateway_event"], ) -def test_lifespan_state(mock_aws_api_gateway_event, lifespan_state, lifespan) -> None: +def test_lifespan_state(mock_aws_api_gateway_event, lifespan: Literal["on", "auto"]) -> None: startup_complete = False shutdown_complete = False - async def app(scope, receive, send): + async def app(scope: Scope, receive: Receive, send: 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) + scope["state"].update({"test_key": b"Hello, world!"}) await send({"type": "lifespan.startup.complete"}) startup_complete = True elif message["type"] == "lifespan.shutdown": @@ -237,7 +236,6 @@ async def app(scope, receive, send): return if scope["type"] == "http": - assert lifespan_state.items() <= scope["state"].items() await send( { "type": "http.response.start", @@ -245,7 +243,7 @@ async def app(scope, receive, send): "headers": [[b"content-type", b"text/plain; charset=utf-8"]], } ) - await send({"type": "http.response.body", "body": b"Hello, world!"}) + await send({"type": "http.response.body", "body": scope["state"]["test_key"]}) handler = Mangum(app, lifespan=lifespan) response = handler(mock_aws_api_gateway_event, {}) From c9df4307425ba1a5f1528bbf6106ecac5341ed72 Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Thu, 26 Sep 2024 22:28:05 +0200 Subject: [PATCH 8/8] Use typing_extensions.Literal --- tests/test_lifespan.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_lifespan.py b/tests/test_lifespan.py index ef1de7d..c5ab74d 100644 --- a/tests/test_lifespan.py +++ b/tests/test_lifespan.py @@ -1,10 +1,10 @@ import logging -from typing import Literal import pytest from quart import Quart from starlette.applications import Starlette from starlette.responses import PlainTextResponse +from typing_extensions import Literal from mangum import Mangum from mangum.exceptions import LifespanFailure