New shutdown() method added to the client (Fixes miguelgrinberg#1333)

miguelgrinberg · miguelgrinberg · commit 811e044a46b7 · 2024-05-19T20:30:59.000+01:00
diff --git a/src/socketio/async_client.py b/src/socketio/async_client.py
@@ -318,6 +318,21 @@ async def disconnect(self):
                                     namespace=n))
         await self.eio.disconnect(abort=True)
 
+    async def shutdown(self):
+        """Stop the client.
+
+        If the client is connected to a server, it is disconnected. If the
+        client is attempting to reconnect to server, the reconnection attempts
+        are stopped. If the client is not connected to a server and is not
+        attempting to reconnect, then this function does nothing.
+        """
+        if self.connected:
+            await self.disconnect()
+        elif self._reconnect_task:  # pragma: no branch
+            self._reconnect_abort.set()
+            print(self._reconnect_task)
+            await self._reconnect_task
+
     def start_background_task(self, target, *args, **kwargs):
         """Start a background task using the appropriate async model.
 
@@ -467,15 +482,20 @@ async def _handle_reconnect(self):
             self.logger.info(
                 'Connection failed, new attempt in {:.02f} seconds'.format(
                     delay))
+            abort = False
             try:
                 await asyncio.wait_for(self._reconnect_abort.wait(), delay)
+                abort = True
+            except asyncio.TimeoutError:
+                pass
+            except asyncio.CancelledError:  # pragma: no cover
+                abort = True
+            if abort:
                 self.logger.info('Reconnect task aborted')
                 for n in self.connection_namespaces:
                     await self._trigger_event('__disconnect_final',
                                               namespace=n)
                 break
-            except (asyncio.TimeoutError, asyncio.CancelledError):
-                pass
             attempt_count += 1
             try:
                 await self.connect(self.connection_url,
diff --git a/src/socketio/client.py b/src/socketio/client.py
@@ -298,6 +298,20 @@ def disconnect(self):
                 packet.DISCONNECT, namespace=n))
         self.eio.disconnect(abort=True)
 
+    def shutdown(self):
+        """Stop the client.
+
+        If the client is connected to a server, it is disconnected. If the
+        client is attempting to reconnect to server, the reconnection attempts
+        are stopped. If the client is not connected to a server and is not
+        attempting to reconnect, then this function does nothing.
+        """
+        if self.connected:
+            self.disconnect()
+        elif self._reconnect_task:  # pragma: no branch
+            self._reconnect_abort.set()
+            self._reconnect_task.join()
+
     def start_background_task(self, target, *args, **kwargs):
         """Start a background task using the appropriate async model.
 
diff --git a/tests/async/test_client.py b/tests/async/test_client.py
@@ -990,6 +990,66 @@ def test_handle_reconnect_aborted(self, random, wait_for):
         c._trigger_event.mock.assert_called_once_with('__disconnect_final',
                                                       namespace='/')
 
+    def test_shutdown_disconnect(self):
+        c = async_client.AsyncClient()
+        c.connected = True
+        c.namespaces = {'/': '1'}
+        c._trigger_event = AsyncMock()
+        c._send_packet = AsyncMock()
+        c.eio = mock.MagicMock()
+        c.eio.disconnect = AsyncMock()
+        c.eio.state = 'connected'
+        _run(c.shutdown())
+        assert c._trigger_event.mock.call_count == 0
+        assert c._send_packet.mock.call_count == 1
+        expected_packet = packet.Packet(packet.DISCONNECT, namespace='/')
+        assert (
+            c._send_packet.mock.call_args_list[0][0][0].encode()
+            == expected_packet.encode()
+        )
+        c.eio.disconnect.mock.assert_called_once_with(abort=True)
+
+    def test_shutdown_disconnect_namespaces(self):
+        c = async_client.AsyncClient()
+        c.connected = True
+        c.namespaces = {'/foo': '1', '/bar': '2'}
+        c._trigger_event = AsyncMock()
+        c._send_packet = AsyncMock()
+        c.eio = mock.MagicMock()
+        c.eio.disconnect = AsyncMock()
+        c.eio.state = 'connected'
+        _run(c.shutdown())
+        assert c._trigger_event.mock.call_count == 0
+        assert c._send_packet.mock.call_count == 2
+        expected_packet = packet.Packet(packet.DISCONNECT, namespace='/foo')
+        assert (
+            c._send_packet.mock.call_args_list[0][0][0].encode()
+            == expected_packet.encode()
+        )
+        expected_packet = packet.Packet(packet.DISCONNECT, namespace='/bar')
+        assert (
+            c._send_packet.mock.call_args_list[1][0][0].encode()
+            == expected_packet.encode()
+        )
+
+    @mock.patch('socketio.client.random.random', side_effect=[1, 0, 0.5])
+    def test_shutdown_reconnect(self, random):
+        c = async_client.AsyncClient()
+        c.connection_namespaces = ['/']
+        c._reconnect_task = AsyncMock()()
+        c._trigger_event = AsyncMock()
+        c.connect = AsyncMock(side_effect=exceptions.ConnectionError)
+
+        async def r():
+            task = c.start_background_task(c._handle_reconnect)
+            await asyncio.sleep(0.1)
+            await c.shutdown()
+            await task
+
+        _run(r())
+        c._trigger_event.mock.assert_called_once_with('__disconnect_final',
+                                                      namespace='/')
+
     def test_handle_eio_connect(self):
         c = async_client.AsyncClient()
         c.connection_namespaces = ['/', '/foo']
diff --git a/tests/common/test_client.py b/tests/common/test_client.py
@@ -1,4 +1,5 @@
 import logging
+import time
 import unittest
 from unittest import mock
 
@@ -636,6 +637,7 @@ def test_disconnect(self):
 
     def test_disconnect_namespaces(self):
         c = client.Client()
+        c.connected = True
         c.namespaces = {'/foo': '1', '/bar': '2'}
         c._trigger_event = mock.MagicMock()
         c._send_packet = mock.MagicMock()
@@ -1128,6 +1130,61 @@ def test_handle_reconnect_aborted(self, random):
         c._trigger_event.assert_called_once_with('__disconnect_final',
                                                  namespace='/')
 
+    def test_shutdown_disconnect(self):
+        c = client.Client()
+        c.connected = True
+        c.namespaces = {'/': '1'}
+        c._trigger_event = mock.MagicMock()
+        c._send_packet = mock.MagicMock()
+        c.eio = mock.MagicMock()
+        c.eio.state = 'connected'
+        c.shutdown()
+        assert c._trigger_event.call_count == 0
+        assert c._send_packet.call_count == 1
+        expected_packet = packet.Packet(packet.DISCONNECT, namespace='/')
+        assert (
+            c._send_packet.call_args_list[0][0][0].encode()
+            == expected_packet.encode()
+        )
+        c.eio.disconnect.assert_called_once_with(abort=True)
+
+    def test_shutdown_disconnect_namespaces(self):
+        c = client.Client()
+        c.connected = True
+        c.namespaces = {'/foo': '1', '/bar': '2'}
+        c._trigger_event = mock.MagicMock()
+        c._send_packet = mock.MagicMock()
+        c.eio = mock.MagicMock()
+        c.eio.state = 'connected'
+        c.shutdown()
+        assert c._trigger_event.call_count == 0
+        assert c._send_packet.call_count == 2
+        expected_packet = packet.Packet(packet.DISCONNECT, namespace='/foo')
+        assert (
+            c._send_packet.call_args_list[0][0][0].encode()
+            == expected_packet.encode()
+        )
+        expected_packet = packet.Packet(packet.DISCONNECT, namespace='/bar')
+        assert (
+            c._send_packet.call_args_list[1][0][0].encode()
+            == expected_packet.encode()
+        )
+
+    @mock.patch('socketio.client.random.random', side_effect=[1, 0, 0.5])
+    def test_shutdown_reconnect(self, random):
+        c = client.Client()
+        c.connection_namespaces = ['/']
+        c._reconnect_task = mock.MagicMock()
+        c._trigger_event = mock.MagicMock()
+        c.connect = mock.MagicMock(side_effect=exceptions.ConnectionError)
+        task = c.start_background_task(c._handle_reconnect)
+        time.sleep(0.1)
+        c.shutdown()
+        task.join()
+        c._trigger_event.assert_called_once_with('__disconnect_final',
+                                                 namespace='/')
+        assert c._reconnect_task.join.called_once_with()
+
     def test_handle_eio_connect(self):
         c = client.Client()
         c.connection_namespaces = ['/', '/foo']
