Improved signal handler for the async client (Fixes miguelgrinberg/py…

…thon-socketio#523)

engineio/asyncio_client.py

@@ -1,5 +1,7 @@
 import asyncio
+import signal
 import ssl
+import threading
 
 try:
     import aiohttp
@@ -12,6 +14,23 @@
 from . import packet
 from . import payload
 
+async_signal_handler_set = False
+
+
+def async_signal_handler():
+    """SIGINT handler.
+
+    Disconnect all active async clients.
+    """
+    async def _handler():
+        for c in client.connected_clients[:]:
+            if c.is_asyncio_based():
+                await c.disconnect()
+        else:  # pragma: no cover
+            pass
+
+    asyncio.ensure_future(_handler())
+
 
 class AsyncClient(client.Client):
     """An Engine.IO client for asyncio.
@@ -61,6 +80,12 @@ async def connect(self, url, headers={}, transports=None,
             eio = engineio.Client()
             await eio.connect('http://localhost:5000')
         """
+        global async_signal_handler_set
+        if not async_signal_handler_set and \
+                threading.current_thread() == threading.main_thread():
+            asyncio.get_event_loop().add_signal_handler(signal.SIGINT,
+                                                        async_signal_handler)
+            async_signal_handler_set = True
         if self.state != 'disconnected':
             raise ValueError('Client is not in a disconnected state')
         valid_transports = ['polling', 'websocket']

engineio/client.py

@@ -35,10 +35,8 @@ def signal_handler(sig, frame):
     Disconnect all active clients and then invoke the original signal handler.
     """
     for client in connected_clients[:]:
-        if client.is_asyncio_based():
-            client.start_background_task(client.disconnect, abort=True)
-        else:
-            client.disconnect(abort=True)
+        if not client.is_asyncio_based():
+            client.disconnect()
     if callable(original_signal_handler):
         return original_signal_handler(sig, frame)
     else:  # pragma: no cover

tests/asyncio/test_asyncio_client.py

@@ -1510,3 +1510,17 @@ def test_write_loop_websocket_bad_connection(self):
         )
         _run(c._write_loop())
         assert c.state == 'connected'
+
+    @mock.patch('engineio.client.original_signal_handler')
+    def test_signal_handler(self, original_handler):
+        clients = [mock.MagicMock(), mock.MagicMock()]
+        client.connected_clients = clients[:]
+        client.connected_clients[0].is_asyncio_based.return_value = False
+        client.connected_clients[1].is_asyncio_based.return_value = True
+
+        async def test():
+            asyncio_client.async_signal_handler()
+
+        asyncio.get_event_loop().run_until_complete(test())
+        clients[0].disconnect.assert_not_called()
+        clients[1].disconnect.assert_called_once_with()

tests/common/test_client.py

@@ -1486,7 +1486,5 @@ def test_signal_handler(self, original_handler):
         client.connected_clients[0].is_asyncio_based.return_value = False
         client.connected_clients[1].is_asyncio_based.return_value = True
         client.signal_handler('sig', 'frame')
-        clients[0].disconnect.assert_called_once_with(abort=True)
-        clients[1].start_background_task.assert_called_once_with(
-            clients[1].disconnect, abort=True
-        )
+        clients[0].disconnect.assert_called_once_with()
+        clients[1].disconnect.assert_not_called()

tox.ini

@@ -1,5 +1,5 @@
 [tox]
-envlist=flake8,py35,py36,py37,pypy3,docs
+envlist=flake8,py35,py36,py37,py38,pypy3,docs
 skip_missing_interpreters=True
 
 [testenv]
@@ -19,6 +19,7 @@ basepython =
     py35: python3.5
     py36: python3.6
     py37: python3.7
+    py38: python3.8
     pypy3: pypy3
     coverage: python3.7
     docs: python3.7