Automatically detect lifespan support

jordaneremieff · jordaneremieff · commit 88a76cc91e83 · 2020-04-19T19:26:35.000+10:00
diff --git a/README.md b/README.md
@@ -77,7 +77,7 @@ async def app(scope, receive, send):
     await send({"type": "http.response.body", "body": b"Hello, world!"})
 
 
-handler = Mangum(app, enable_lifespan=False) # disable lifespan for raw ASGI example
+handler = Mangum(app)
 ```
 
 ## WebSockets (experimental)
diff --git a/mangum/adapter.py b/mangum/adapter.py
@@ -48,7 +48,7 @@ def __post_init__(self) -> None:
             loop = asyncio.get_event_loop()
             self.lifespan = Lifespan(self.app, logger=self.logger)
             loop.create_task(self.lifespan.run())
-            loop.run_until_complete(self.lifespan.wait_startup())
+            loop.run_until_complete(self.lifespan.startup())
 
     def __call__(self, event: dict, context: dict) -> dict:
         try:
@@ -74,8 +74,9 @@ def handler(self, event: dict, context: dict) -> dict:
             response = self.handle_http(event, context, is_http_api=is_http_api)
 
         if self.enable_lifespan:
-            loop = asyncio.get_event_loop()
-            loop.run_until_complete(self.lifespan.wait_shutdown())
+            if self.lifespan.is_supported:
+                loop = asyncio.get_event_loop()
+                loop.run_until_complete(self.lifespan.shutdown())
 
         return response
 
diff --git a/mangum/exceptions.py b/mangum/exceptions.py
@@ -1,2 +1,6 @@
 class ASGIWebSocketCycleException(Exception):
     """Raise when an exception occurs within an ASGI websocket cycle"""
+
+
+class LifespanFailure(Exception):
+    """Raise when an error occurs in a lifespan event"""
diff --git a/mangum/lifespan.py b/mangum/lifespan.py
@@ -2,56 +2,66 @@
 import asyncio
 from dataclasses import dataclass
 
-from mangum.types import ASGIApp, Message, Send, Receive
+from mangum.types import ASGIApp, Message
+from mangum.exceptions import LifespanFailure
 
 
 @dataclass
 class Lifespan:
     app: ASGIApp
     logger: logging.Logger
-    startup_event: asyncio.Event = asyncio.Event()
-    shutdown_event: asyncio.Event = asyncio.Event()
-    app_queue: asyncio.Queue = asyncio.Queue()
+    is_supported: bool = False
+    has_error: bool = False
+
+    def __post_init__(self) -> None:
+        self.app_queue: asyncio.Queue = asyncio.Queue()
+        self.startup_event: asyncio.Event = asyncio.Event()
+        self.shutdown_event: asyncio.Event = asyncio.Event()
+
+    async def startup(self) -> None:
+        self.logger.info("Waiting for application startup.")
+        if self.is_supported:
+            await self.app_queue.put({"type": "lifespan.startup"})
+            await self.startup_event.wait()
+            if self.has_error:
+                self.logger.error("Application startup failed.")
+            else:
+                self.logger.info("Application startup complete.")
+
+    async def shutdown(self) -> None:
+        if self.has_error:
+            return
+        self.logger.info("Waiting for application shutdown.")
+        await self.app_queue.put({"type": "lifespan.shutdown"})
+        await self.shutdown_event.wait()
 
     async def run(self) -> None:
-        receive, send = (self.receiver(), self.sender())
         try:
-            await self.app({"type": "lifespan"}, receive, send)
-        except BaseException as exc:  # pragma: no cover
-            self.logger.error(f"Exception in 'lifespan' protocol: {exc}")
-        finally:
+            await self.app({"type": "lifespan"}, self.receive, self.send)
+        except BaseException as exc:
             self.startup_event.set()
             self.shutdown_event.set()
+            self.has_error = True
+            if not self.is_supported:
+                self.logger.info("ASGI 'lifespan' protocol appears unsupported.")
+            else:
+                self.logger.error("Exception in 'lifespan' protocol.", exc_info=exc)
 
-    def sender(self) -> Send:
-        # startup_event, shutdown_event = self.startup_event, self.shutdown_event
-
-        async def send(message: Message) -> None:
-            message_type = message["type"]
-            if message_type == "lifespan.startup.complete":
-                self.startup_event.set()
-            elif message_type == "lifespan.shutdown.complete":
-                self.shutdown_event.set()
-            else:  # pragma: no cover
-                raise RuntimeError(
-                    f"Expected lifespan message type, received: {message_type}"
-                )
-            return None
+    async def receive(self) -> Message:
+        self.is_supported = True
 
-        return send
+        return await self.app_queue.get()
 
-    def receiver(self) -> Receive:
-        async def receive() -> Message:
-            return await self.app_queue.get()
+    async def send(self, message: Message) -> None:
+        if not self.is_supported:
+            raise LifespanFailure("Lifespan unsupported.")
 
-        return receive
-
-    async def wait_startup(self) -> None:
-        self.logger.info("Waiting for application startup.")
-        await self.app_queue.put({"type": "lifespan.startup"})
-        await self.startup_event.wait()
-
-    async def wait_shutdown(self) -> None:
-        self.logger.info("Waiting for application shutdown.")
-        await self.app_queue.put({"type": "lifespan.shutdown"})
-        await self.shutdown_event.wait()
+        message_type = message["type"]
+        if message_type == "lifespan.startup.complete":
+            self.startup_event.set()
+        elif message_type == "lifespan.shutdown.complete":
+            self.shutdown_event.set()
+        else:  # pragma: no cover
+            raise RuntimeError(
+                f"Expected lifespan message type, received: {message_type}"
+            )
diff --git a/tests/test_lifespan.py b/tests/test_lifespan.py
@@ -4,7 +4,6 @@
 from starlette.applications import Starlette
 from starlette.responses import PlainTextResponse
 
-
 from mangum import Mangum
 
 # One (or more) of Quart's dependencies does not support Python 3.8, ignore this case.
@@ -18,9 +17,156 @@
 
 
 @pytest.mark.parametrize("mock_http_event", [["GET", None, None]], indirect=True)
-def test_starlette_response(mock_http_event) -> None:
+def test_lifespan_startup_error(mock_http_event) -> None:
+    async def app(scope, receive, send):
+        if scope["type"] == "lifespan":
+            while True:
+                message = await receive()
+                if message["type"] == "lifespan.startup":
+                    raise Exception("error")
+        else:
+            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)
+    assert handler.lifespan.is_supported
+    assert handler.lifespan.has_error
+
+    response = handler(mock_http_event, {})
+    assert response == {
+        "statusCode": 200,
+        "isBase64Encoded": False,
+        "headers": {"content-type": "text/plain; charset=utf-8"},
+        "body": "Hello, world!",
+    }
+
+
+@pytest.mark.parametrize("mock_http_event", [["GET", None, None]], indirect=True)
+def test_lifespan(mock_http_event) -> 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":
+                    await send({"type": "lifespan.startup.complete"})
+                    startup_complete = True
+                elif message["type"] == "lifespan.shutdown":
+                    await send({"type": "lifespan.shutdown.complete"})
+                    shutdown_complete = True
+                    return
+        else:
+            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)
+    assert startup_complete
+
+    response = handler(mock_http_event, {})
+    assert shutdown_complete
+    assert response == {
+        "statusCode": 200,
+        "isBase64Encoded": False,
+        "headers": {"content-type": "text/plain; charset=utf-8"},
+        "body": "Hello, world!",
+    }
+
+
+@pytest.mark.parametrize("mock_http_event", [["GET", None, None]], indirect=True)
+def test_lifespan_unsupported(mock_http_event) -> None:
+    async def app(scope, receive, send):
+        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)
+    assert not handler.lifespan.is_supported
+
+    response = handler(mock_http_event, {})
+    assert response == {
+        "statusCode": 200,
+        "isBase64Encoded": False,
+        "headers": {"content-type": "text/plain; charset=utf-8"},
+        "body": "Hello, world!",
+    }
+
+
+@pytest.mark.parametrize("mock_http_event", [["GET", None, None]], indirect=True)
+def test_lifespan_disabled(mock_http_event) -> None:
+    async def app(scope, receive, send):
+        assert scope["type"] == "http"
+        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, enable_lifespan=False)
+
+    response = handler(mock_http_event, {})
+    assert response == {
+        "statusCode": 200,
+        "isBase64Encoded": False,
+        "headers": {"content-type": "text/plain; charset=utf-8"},
+        "body": "Hello, world!",
+    }
+
+
+@pytest.mark.parametrize("mock_http_event", [["GET", None, None]], indirect=True)
+def test_lifespan_supported_with_error(mock_http_event) -> None:
+    async def app(scope, receive, send):
+        await receive()
+        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)
+    assert handler.lifespan.is_supported
+    assert handler.lifespan.has_error
+
+    response = handler(mock_http_event, {})
+    assert response == {
+        "statusCode": 200,
+        "isBase64Encoded": False,
+        "headers": {"content-type": "text/plain; charset=utf-8"},
+        "body": "Hello, world!",
+    }
+
+
+@pytest.mark.parametrize("mock_http_event", [["GET", None, None]], indirect=True)
+def test_starlette_lifespan(mock_http_event) -> None:
+    startup_complete = False
+    shutdown_complete = False
+
     path = mock_http_event["path"]
     app = Starlette()
 
@@ -67,7 +213,7 @@ def homepage(request):
 )
 @pytest.mark.skipif(IS_PY36, reason="Quart does not support Python 3.6.")
 @pytest.mark.parametrize("mock_http_event", [["GET", None, None]], indirect=True)
-def test_quart_app(mock_http_event) -> None:
+def test_quart_lifespan(mock_http_event) -> None:
     startup_complete = False
     shutdown_complete = False
     path = mock_http_event["path"]
