Better handling of connection state

miguelgrinberg · Sep 26, 2015 · c3715e6 · c3715e6
1 parent 48d999d
commit c3715e6
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 7 deletions.
diff --git a/engineio/server.py b/engineio/server.py
@@ -266,7 +266,7 @@ def _handle_connect(self, environ, start_response, transport, b64=False):
 
         pkt = packet.Packet(
             packet.OPEN, {'sid': sid,
-                          'upgrades': self._upgrades(sid),
+                          'upgrades': self._upgrades(sid, transport),
                           'pingTimeout': int(self.ping_timeout * 1000),
                           'pingInterval': int(self.ping_interval * 1000)})
         s.send(pkt)
@@ -280,15 +280,17 @@ def _handle_connect(self, environ, start_response, transport, b64=False):
             s.handle_get_request(environ, start_response)
             return self._ok()
         else:
+            s.connected = True
             headers = None
             if self.cookie:
                 headers = [('Set-Cookie', self.cookie + '=' + sid)]
             return self._ok(s.poll(), headers=headers, b64=b64)
 
-    def _upgrades(self, sid):
+    def _upgrades(self, sid, transport):
         """Return the list of possible upgrades for a client connection."""
         if not self.allow_upgrades or self._get_socket(sid).upgraded or \
-                self.async['websocket_class'] is None:
+                self.async['websocket_class'] is None or \
+                transport == 'websocket':
             return []
         return ['websocket']
 

diff --git a/engineio/socket.py b/engineio/socket.py
@@ -73,7 +73,6 @@ def handle_get_request(self, environ, start_response):
                                     self.sid, transport)
             return getattr(self, '_upgrade_' + transport)(environ,
                                                           start_response)
-        self.connected = True
         try:
             packets = self.poll()
         except IOError as e:

diff --git a/tests/test_server.py b/tests/test_server.py
@@ -197,12 +197,15 @@ def test_close_all_sockets(self):
     def test_upgrades(self):
         s = server.Server()
         s.sockets['foo'] = self._get_mock_socket()
-        self.assertEqual(s._upgrades('foo'), ['websocket'])
+        self.assertEqual(s._upgrades('foo', 'polling'), ['websocket'])
+        self.assertEqual(s._upgrades('foo', 'websocket'), [])
         s.sockets['foo'].upgraded = True
-        self.assertEqual(s._upgrades('foo'), [])
+        self.assertEqual(s._upgrades('foo', 'polling'), [])
+        self.assertEqual(s._upgrades('foo', 'websocket'), [])
         s.allow_upgrades = False
         s.sockets['foo'].upgraded = True
-        self.assertEqual(s._upgrades('foo'), [])
+        self.assertEqual(s._upgrades('foo', 'polling'), [])
+        self.assertEqual(s._upgrades('foo', 'websocket'), [])
 
     def test_transport(self):
         s = server.Server()