Fixed disconnect implementation when using a message queue (Fixes miguelgrinberg#1002)

miguelgrinberg · miguelgrinberg · commit f56ef6f0401b · 2022-08-15T16:01:52.000+01:00
diff --git a/src/socketio/base_manager.py b/src/socketio/base_manager.py
@@ -94,7 +94,7 @@ def pre_disconnect(self, sid, namespace):
         self.pending_disconnect[namespace].append(sid)
         return self.rooms[namespace][None].get(sid)
 
-    def disconnect(self, sid, namespace):
+    def disconnect(self, sid, namespace, **kwargs):
         """Register a client disconnect from a namespace."""
         if namespace not in self.rooms:
             return
diff --git a/src/socketio/pubsub_manager.py b/src/socketio/pubsub_manager.py
@@ -75,7 +75,10 @@ def can_disconnect(self, sid, namespace):
             self._publish({'method': 'disconnect', 'sid': sid,
                            'namespace': namespace or '/'})
 
-    def disconnect(self, sid, namespace=None):
+    def disconnect(self, sid, namespace=None, **kwargs):
+        if kwargs.get('ignore_queue'):
+            return super(PubSubManager, self).disconnect(
+                sid, namespace=namespace)
         self._publish({'method': 'disconnect', 'sid': sid,
                        'namespace': namespace or '/'})
 
diff --git a/src/socketio/server.py b/src/socketio/server.py
@@ -565,7 +565,8 @@ def disconnect(self, sid, namespace=None, ignore_queue=False):
             self._send_packet(eio_sid, self.packet_class(
                 packet.DISCONNECT, namespace=namespace))
             self._trigger_event('disconnect', namespace, sid)
-            self.manager.disconnect(sid, namespace=namespace)
+            self.manager.disconnect(sid, namespace=namespace,
+                                    ignore_queue=True)
 
     def transport(self, sid):
         """Return the name of the transport used by the client.
@@ -693,7 +694,7 @@ def _handle_connect(self, eio_sid, namespace, data):
                 self._send_packet(eio_sid, self.packet_class(
                     packet.CONNECT_ERROR, data=fail_reason,
                     namespace=namespace))
-            self.manager.disconnect(sid, namespace)
+            self.manager.disconnect(sid, namespace, ignore_queue=True)
         elif not self.always_connect:
             self._send_packet(eio_sid, self.packet_class(
                 packet.CONNECT, {'sid': sid}, namespace=namespace))
@@ -706,7 +707,7 @@ def _handle_disconnect(self, eio_sid, namespace):
             return
         self.manager.pre_disconnect(sid, namespace=namespace)
         self._trigger_event('disconnect', namespace, sid)
-        self.manager.disconnect(sid, namespace)
+        self.manager.disconnect(sid, namespace, ignore_queue=True)
 
     def _handle_event(self, eio_sid, namespace, id, data):
         """Handle an incoming client event."""
diff --git a/tests/common/test_pubsub_manager.py b/tests/common/test_pubsub_manager.py
@@ -175,6 +175,13 @@ def test_disconnect(self):
             {'method': 'disconnect', 'sid': 'foo', 'namespace': '/'}
         )
 
+    def test_disconnect_ignore_queue(self):
+        sid = self.pm.connect('123', '/')
+        self.pm.pre_disconnect(sid, '/')
+        self.pm.disconnect(sid, ignore_queue=True)
+        self.pm._publish.assert_not_called()
+        assert not self.pm.is_connected(sid, '/')
+
     def test_close_room(self):
         self.pm.close_room('foo')
         self.pm._publish.assert_called_once_with(
diff --git a/tests/common/test_server.py b/tests/common/test_server.py
@@ -532,7 +532,8 @@ def test_handle_disconnect(self, eio):
         s._handle_eio_message('123', '0')
         s._handle_eio_disconnect('123')
         handler.assert_called_once_with('1')
-        s.manager.disconnect.assert_called_once_with('1', '/')
+        s.manager.disconnect.assert_called_once_with('1', '/',
+                                                     ignore_queue=True)
         assert s.environ == {}
 
     def test_handle_disconnect_namespace(self, eio):
