From c419fc5481846ab026ca847234f5ebe5420e510a Mon Sep 17 00:00:00 2001
From: Miguel Grinberg <miguel.grinberg@gmail.com>
Date: Thu, 14 Sep 2023 00:08:02 +0100
Subject: [PATCH] Add a shutdown() function for the server

---
 src/socketio/asyncio_server.py       | 9 +++++++++
 src/socketio/server.py               | 9 +++++++++
 tests/asyncio/test_asyncio_server.py | 6 ++++++
 tests/common/test_server.py          | 5 +++++
 4 files changed, 29 insertions(+)

diff --git a/src/socketio/asyncio_server.py b/src/socketio/asyncio_server.py
index aff8812a..acdb265d 100644
--- a/src/socketio/asyncio_server.py
+++ b/src/socketio/asyncio_server.py
@@ -387,6 +387,15 @@ async def disconnect(self, sid, namespace=None, ignore_queue=False):
             await self.manager.disconnect(sid, namespace=namespace,
                                           ignore_queue=True)
 
+    async def shutdown(self):
+        """Stop Socket.IO background tasks.
+
+        This method stops all background activity initiated by the Socket.IO
+        server. It must be called before shutting down the web server.
+        """
+        self.logger.info('Socket.IO is shutting down')
+        await self.eio.shutdown()
+
     async def handle_request(self, *args, **kwargs):
         """Handle an HTTP request from the client.
 
diff --git a/src/socketio/server.py b/src/socketio/server.py
index 48d55983..8675b1e2 100644
--- a/src/socketio/server.py
+++ b/src/socketio/server.py
@@ -570,6 +570,15 @@ def disconnect(self, sid, namespace=None, ignore_queue=False):
             self.manager.disconnect(sid, namespace=namespace,
                                     ignore_queue=True)
 
+    def shutdown(self):
+        """Stop Socket.IO background tasks.
+
+        This method stops all background activity initiated by the Socket.IO
+        server. It must be called before shutting down the web server.
+        """
+        self.logger.info('Socket.IO is shutting down')
+        self.eio.shutdown()
+
     def transport(self, sid):
         """Return the name of the transport used by the client.
 
diff --git a/tests/asyncio/test_asyncio_server.py b/tests/asyncio/test_asyncio_server.py
index 0adce679..ef9c2bb6 100644
--- a/tests/asyncio/test_asyncio_server.py
+++ b/tests/asyncio/test_asyncio_server.py
@@ -981,6 +981,12 @@ def test_async_handlers(self, eio):
             None,
         )
 
+    def test_shutdown(self, eio):
+        s = asyncio_server.AsyncServer()
+        s.eio.shutdown = AsyncMock()
+        _run(s.shutdown())
+        s.eio.shutdown.mock.assert_called_once_with()
+
     def test_start_background_task(self, eio):
         s = asyncio_server.AsyncServer()
         s.start_background_task('foo', 'bar', baz='baz')
diff --git a/tests/common/test_server.py b/tests/common/test_server.py
index 9285e570..08c59ac8 100644
--- a/tests/common/test_server.py
+++ b/tests/common/test_server.py
@@ -917,6 +917,11 @@ def test_async_handlers(self, eio):
             None,
         )
 
+    def test_shutdown(self, eio):
+        s = server.Server()
+        s.shutdown()
+        s.eio.shutdown.assert_called_once_with()
+
     def test_start_background_task(self, eio):
         s = server.Server()
         s.start_background_task('foo', 'bar', baz='baz')