Emit ASGI lifespan shutdown event (Fixes #200)

miguelgrinberg · Nov 27, 2020 · 2fb1774 · 2fb1774
1 parent 00330bb
commit 2fb1774
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 30 deletions.
diff --git a/engineio/async_drivers/asgi.py b/engineio/async_drivers/asgi.py
@@ -81,27 +81,29 @@ async def serve_static_file(self, static_file, receive,
                 await self.not_found(receive, send)
 
     async def lifespan(self, receive, send):
-        event = await receive()
-        if event['type'] == 'lifespan.startup':
-            if self.on_startup:
-                try:
-                    await self.on_startup() \
-                        if asyncio.iscoroutinefunction(self.on_startup) else \
-                        self.on_startup()
-                except:
-                    await send({'type': 'lifespan.startup.failed'})
-                    return
-            await send({'type': 'lifespan.startup.complete'})
-        elif event['type'] == 'lifespan.shutdown':
-            if self.on_shutdown:
-                try:
-                    await self.on_shutdown() \
-                        if asyncio.iscoroutinefunction(self.on_shutdown) \
-                        else self.on_shutdown()
-                except:
-                    await send({'type': 'lifespan.shutdown.failed'})
-                    return
-            await send({'type': 'lifespan.shutdown.complete'})
+        while True:
+            event = await receive()
+            if event['type'] == 'lifespan.startup':
+                if self.on_startup:
+                    try:
+                        await self.on_startup() \
+                            if asyncio.iscoroutinefunction(self.on_startup) \
+                            else self.on_startup()
+                    except:
+                        await send({'type': 'lifespan.startup.failed'})
+                        return
+                await send({'type': 'lifespan.startup.complete'})
+            elif event['type'] == 'lifespan.shutdown':
+                if self.on_shutdown:
+                    try:
+                        await self.on_shutdown() \
+                            if asyncio.iscoroutinefunction(self.on_shutdown) \
+                            else self.on_shutdown()
+                    except:
+                        await send({'type': 'lifespan.shutdown.failed'})
+                        return
+                await send({'type': 'lifespan.shutdown.complete'})
+                return
 
     async def not_found(self, receive, send):
         """Return a 404 Not Found error to the client."""

diff --git a/tests/asyncio/test_async_asgi.py b/tests/asyncio/test_async_asgi.py
@@ -125,10 +125,11 @@ def check_path(path, status_code, content_type, body):
     def test_lifespan_startup(self):
         app = async_asgi.ASGIApp('eio')
         scope = {'type': 'lifespan'}
-        receive = AsyncMock(return_value={'type': 'lifespan.startup'})
+        receive = AsyncMock(side_effect=[{'type': 'lifespan.startup'},
+                                         {'type': 'lifespan.shutdown'}])
         send = AsyncMock()
         _run(app(scope, receive, send))
-        send.mock.assert_called_once_with(
+        send.mock.assert_any_call(
             {'type': 'lifespan.startup.complete'}
         )
 
@@ -141,10 +142,11 @@ def startup():
 
         app = async_asgi.ASGIApp('eio', on_startup=startup)
         scope = {'type': 'lifespan'}
-        receive = AsyncMock(return_value={'type': 'lifespan.startup'})
+        receive = AsyncMock(side_effect=[{'type': 'lifespan.startup'},
+                                         {'type': 'lifespan.shutdown'}])
         send = AsyncMock()
         _run(app(scope, receive, send))
-        send.mock.assert_called_once_with(
+        send.mock.assert_any_call(
             {'type': 'lifespan.startup.complete'}
         )
         assert up
@@ -158,10 +160,11 @@ async def startup():
 
         app = async_asgi.ASGIApp('eio', on_startup=startup)
         scope = {'type': 'lifespan'}
-        receive = AsyncMock(return_value={'type': 'lifespan.startup'})
+        receive = AsyncMock(side_effect=[{'type': 'lifespan.startup'},
+                                         {'type': 'lifespan.shutdown'}])
         send = AsyncMock()
         _run(app(scope, receive, send))
-        send.mock.assert_called_once_with(
+        send.mock.assert_any_call(
             {'type': 'lifespan.startup.complete'}
         )
         assert up
@@ -174,7 +177,7 @@ def startup():
 
         app = async_asgi.ASGIApp('eio', on_startup=startup)
         scope = {'type': 'lifespan'}
-        receive = AsyncMock(return_value={'type': 'lifespan.startup'})
+        receive = AsyncMock(side_effect=[{'type': 'lifespan.startup'}])
         send = AsyncMock()
         _run(app(scope, receive, send))
         send.mock.assert_called_once_with({'type': 'lifespan.startup.failed'})
@@ -241,10 +244,13 @@ def shutdown():
     def test_lifespan_invalid(self):
         app = async_asgi.ASGIApp('eio')
         scope = {'type': 'lifespan'}
-        receive = AsyncMock(return_value={'type': 'lifespan.foo'})
+        receive = AsyncMock(side_effect=[{'type': 'lifespan.foo'},
+                                         {'type': 'lifespan.shutdown'}])
         send = AsyncMock()
         _run(app(scope, receive, send))
-        send.mock.assert_not_called()
+        send.mock.assert_called_once_with(
+            {'type': 'lifespan.shutdown.complete'}
+        )
 
     def test_not_found(self):
         app = async_asgi.ASGIApp('eio')