Add ASGI lifespan events (Fixes #322)

docs/extensions.rst

@@ -722,6 +722,18 @@ runs the web application using the Uvicorn web server::
 When using the ASGI support, the ``scope`` dictionary provided by the web
 server is available to request handlers as ``request.asgi_scope``.
 
+The application instance can be initialized with ``lifespan_startup`` and
+``lifespan_shutdown`` arguments, which are invoked when the web server sends
+the ASGI lifespan signals with the ASGI scope as only argument::
+
+    async def startup(scope):
+        pass
+
+    async def shutdown(scope):
+        pass
+
+    app = Microdot(lifespan_startup=startup, lifespan_shutdown=shutdown)
+
 Using a WSGI Web Server
 ^^^^^^^^^^^^^^^^^^^^^^^
 

src/microdot/asgi.py

@@ -48,15 +48,47 @@ class Microdot(BaseMicrodot):
     """A subclass of the core :class:`Microdot <microdot.Microdot>` class that
     implements the ASGI protocol.
 
+    :param startup: An optional function to handle the `lifespan.startup` ASGI
+                    signal.
+    :param shutdown: An optional function to handle the `lifespan.shutdown`
+                     ASGI signal.
+
     This class must be used as the application instance when running under an
     ASGI web server.
     """
-    def __init__(self):
+    def __init__(self, lifespan_startup=None, lifespan_shutdown=None):
         super().__init__()
+        self.lifespan_startup = lifespan_startup
+        self.lifespan_shutdown = lifespan_shutdown
         self.embedded_server = False
 
+    async def handle_lifespan(self, scope, receive, send):
+        while True:
+            message = await receive()
+            if message['type'] == 'lifespan.startup':
+                try:
+                    if self.lifespan_startup:
+                        await self.lifespan_startup(scope)
+                except Exception as e:
+                    await send({'type': 'lifespan.startup.failed',
+                                'message': repr(e)})
+                else:
+                    await send({'type': 'lifespan.startup.complete'})
+            elif message['type'] == 'lifespan.shutdown':  # pragma: no branch
+                try:
+                    if self.lifespan_shutdown:
+                        await self.lifespan_shutdown(scope)
+                except Exception as e:
+                    await send({'type': 'lifespan.shutdown.failed',
+                                'message': repr(e)})
+                else:
+                    await send({'type': 'lifespan.shutdown.complete'})
+                break
+
     async def asgi_app(self, scope, receive, send):
         """An ASGI application."""
+        if scope['type'] == 'lifespan':
+            return await self.handle_lifespan(scope, receive, send)
         if scope['type'] not in ['http', 'websocket']:  # pragma: no cover
             return
         path = scope['path']

tests/test_asgi.py

@@ -170,3 +170,103 @@ class TestASGI(unittest.TestCase):
             self._run(app(scope, receive, send))
 
         kill.assert_called()
+
+    def test_no_lifespan(self):
+        app = Microdot()
+
+        scope = {
+            'type': 'lifespan',
+            'asgi': {'version': '3.0'},
+            'state': {},
+        }
+
+        messages = [
+            {'type': 'lifespan.startup'},
+            {'type': 'lifespan.shutdown'},
+        ]
+        message_iter = iter(messages)
+        sends = []
+
+        async def receive():
+            return next(message_iter)
+
+        async def send(packet):
+            sends.append(packet)
+
+        self._run(app(scope, receive, send))
+        self.assertEqual(sends, [
+            {'type': 'lifespan.startup.complete'},
+            {'type': 'lifespan.shutdown.complete'},
+        ])
+
+    def test_lifespan(self):
+        async def startup(scope):
+            scope['state']['foo'] = 'bar'
+
+        async def shutdown(scope):
+            self.assertEqual(scope['state']['foo'], 'bar')
+            scope['state']['foo'] = 'baz'
+
+        app = Microdot(lifespan_startup=startup, lifespan_shutdown=shutdown)
+
+        scope = {
+            'type': 'lifespan',
+            'asgi': {'version': '3.0'},
+            'state': {},
+        }
+
+        messages = [
+            {'type': 'lifespan.startup'},
+            {'type': 'lifespan.shutdown'},
+        ]
+        message_iter = iter(messages)
+        sends = []
+
+        async def receive():
+            return next(message_iter)
+
+        async def send(packet):
+            sends.append(packet)
+
+        self._run(app(scope, receive, send))
+        self.assertEqual(scope['state']['foo'], 'baz')
+        self.assertEqual(sends, [
+            {'type': 'lifespan.startup.complete'},
+            {'type': 'lifespan.shutdown.complete'},
+        ])
+
+    def test_lifespan_errors(self):
+        async def startup(scope):
+            return scope['scope']['foo']  # KeyError
+
+        async def shutdown(scope):
+            return 1 / 0
+
+        app = Microdot(lifespan_startup=startup, lifespan_shutdown=shutdown)
+
+        scope = {
+            'type': 'lifespan',
+            'asgi': {'version': '3.0'},
+            'state': {},
+        }
+
+        messages = [
+            {'type': 'lifespan.startup'},
+            {'type': 'lifespan.shutdown'},
+        ]
+        message_iter = iter(messages)
+        sends = []
+
+        async def receive():
+            return next(message_iter)
+
+        async def send(packet):
+            sends.append(packet)
+
+        self._run(app(scope, receive, send))
+        self.assertEqual(sends, [
+            {'type': 'lifespan.startup.failed',
+             'message': "KeyError('scope')"},
+            {'type': 'lifespan.shutdown.failed',
+             'message': "ZeroDivisionError('division by zero')"},
+        ])