Fix is_connected() method in test client (Fixes miguelgrinberg/python…

…-socketio#385)
miguelgrinberg · Nov 26, 2019 · a05ff51 · a05ff51
1 parent 1bd15e0
commit a05ff51
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 7 deletions.
diff --git a/flask_socketio/test_client.py b/flask_socketio/test_client.py
@@ -45,7 +45,7 @@ def _mock_send_packet(sid, pkt):
                     pkt.packet_type == packet.BINARY_ACK:
                 self.acks[sid] = {'args': pkt.data,
                                   'namespace': pkt.namespace or '/'}
-            elif pkt.packet_type == packet.DISCONNECT:
+            elif pkt.packet_type in [packet.DISCONNECT, packet.ERROR]:
                 self.connected[pkt.namespace or '/'] = False
 
         self.app = app
@@ -101,16 +101,13 @@ def connect(self, namespace=None, query_string=None, headers=None):
             # inject cookies from Flask
             self.flask_test_client.cookie_jar.inject_wsgi(environ)
         self.connected['/'] = True
-        if self.socketio.server._handle_eio_connect(
-                self.sid, environ) is False:
-            del self.connected['/']
+        self.socketio.server._handle_eio_connect(self.sid, environ)
         if namespace is not None and namespace != '/':
             self.connected[namespace] = True
             pkt = packet.Packet(packet.CONNECT, namespace=namespace)
             with self.app.app_context():
-                if self.socketio.server._handle_eio_message(
-                        self.sid, pkt.encode()) is False:
-                    del self.connected[namespace]
+                self.socketio.server._handle_eio_message(self.sid,
+                                                         pkt.encode())
 
     def disconnect(self, namespace=None):
         """Disconnect the client.

diff --git a/test_socketio.py b/test_socketio.py
@@ -17,6 +17,8 @@
 
 @socketio.on('connect')
 def on_connect():
+    if request.args.get('fail'):
+        return False
     send('connected')
     send(json.dumps(request.args.to_dict(flat=False)))
     send(json.dumps({h: request.headers[h] for h in request.headers.keys()
@@ -300,6 +302,10 @@ def test_connect_namespace_query_string_and_headers(self):
         self.assertEqual(received[2]['args'], '{"My-Custom-Header": "Value"}')
         client.disconnect(namespace='/test')
 
+    def test_connect_rejected(self):
+        client = socketio.test_client(app, query_string='fail=1')
+        self.assertFalse(client.is_connected())
+
     def test_disconnect(self):
         global disconnected
         disconnected = None