Improvements in the connection rejected flow

miguelgrinberg · Jul 29, 2023 · 8051fc4 · 8051fc4
1 parent e72cb6a
commit 8051fc4
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 0 deletions.
diff --git a/src/engineio/asyncio_socket.py b/src/engineio/asyncio_socket.py
@@ -204,6 +204,8 @@ async def writer():
                         await ws.send(pkt.encode())
                 except:
                     break
+            await ws.close()
+
         writer_task = asyncio.ensure_future(writer())
 
         self.server.logger.info(

diff --git a/src/engineio/socket.py b/src/engineio/socket.py
@@ -220,6 +220,8 @@ def writer():
                         ws.send(pkt.encode())
                 except:
                     break
+            ws.close()
+
         writer_task = self.server.start_background_task(writer)
 
         self.server.logger.info(

diff --git a/tests/asyncio/test_asyncio_socket.py b/tests/asyncio/test_asyncio_socket.py
@@ -339,6 +339,7 @@ def test_websocket_read_write(self):
             packet.Packet(packet.MESSAGE, data=foo).encode(),
             None,
         ]
+        ws.close = AsyncMock()
         _run(s._websocket_handler(ws))
         assert s.connected
         assert s.upgraded
@@ -350,6 +351,7 @@ def test_websocket_read_write(self):
             ]
         )
         ws.send.mock.assert_called_with('4bar')
+        ws.close.mock.assert_called()
 
     def test_websocket_upgrade_read_write(self):
         mock_server = self._get_mock_server()
@@ -374,6 +376,7 @@ def test_websocket_upgrade_read_write(self):
             packet.Packet(packet.MESSAGE, data=foo).encode(),
             None,
         ]
+        ws.close = AsyncMock()
         _run(s._websocket_handler(ws))
         assert s.upgraded
         assert mock_server._trigger_event.mock.call_count == 2
@@ -384,6 +387,7 @@ def test_websocket_upgrade_read_write(self):
             ]
         )
         ws.send.mock.assert_called_with('4bar')
+        ws.close.mock.assert_called()
 
     def test_websocket_upgrade_with_payload(self):
         mock_server = self._get_mock_server()
@@ -398,6 +402,7 @@ def test_websocket_upgrade_with_payload(self):
             packet.Packet(packet.PING, data=probe).encode(),
             packet.Packet(packet.UPGRADE, data='2').encode(),
         ]
+        ws.close = AsyncMock()
         _run(s._websocket_handler(ws))
         assert s.upgraded
 
@@ -415,6 +420,7 @@ def test_websocket_upgrade_with_backlog(self):
             packet.Packet(packet.PING, data=probe).encode(),
             packet.Packet(packet.UPGRADE, data='2').encode(),
         ]
+        ws.close = AsyncMock()
         s.upgrading = True
         _run(s.send(packet.Packet(packet.MESSAGE, data=foo)))
         environ = {'REQUEST_METHOD': 'GET', 'QUERY_STRING': 'sid=sid'}
@@ -454,6 +460,7 @@ def test_websocket_read_write_wait_fail(self):
             RuntimeError,
         ]
         ws.send.mock.side_effect = [None, RuntimeError]
+        ws.close = AsyncMock()
         _run(s._websocket_handler(ws))
         assert s.closed
 
@@ -495,6 +502,7 @@ def test_websocket_ignore_invalid_packet(self):
             packet.Packet(packet.MESSAGE, data=foo).encode(),
             None,
         ]
+        ws.close = AsyncMock()
         _run(s._websocket_handler(ws))
         assert s.connected
         assert mock_server._trigger_event.mock.call_count == 2
@@ -505,6 +513,7 @@ def test_websocket_ignore_invalid_packet(self):
             ]
         )
         ws.send.mock.assert_called_with('4bar')
+        ws.close.mock.assert_called()
 
     def test_send_after_close(self):
         mock_server = self._get_mock_server()