Skip to content

Commit

Permalink
Use request timeout when making a WebSocket connection (#286)
Browse files Browse the repository at this point in the history
  • Loading branch information
jpfarias committed Jul 19, 2022
1 parent cb2d7a6 commit f6df30b
Show file tree
Hide file tree
Showing 4 changed files with 26 additions and 3 deletions.
5 changes: 3 additions & 2 deletions src/engineio/asyncio_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -304,11 +304,12 @@ async def _connect_websocket(self, url, headers, engineio_path):
ssl_context.verify_mode = ssl.CERT_NONE
ws = await self.http.ws_connect(
websocket_url + self._get_url_timestamp(),
headers=headers, ssl=ssl_context)
headers=headers, ssl=ssl_context,
timeout=self.request_timeout)
else:
ws = await self.http.ws_connect(
websocket_url + self._get_url_timestamp(),
headers=headers)
headers=headers, timeout=self.request_timeout)
except (aiohttp.client_exceptions.WSServerHandshakeError,
aiohttp.client_exceptions.ServerConnectionError,
aiohttp.client_exceptions.ClientConnectionError):
Expand Down
3 changes: 2 additions & 1 deletion src/engineio/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -417,7 +417,8 @@ def _connect_websocket(self, url, headers, engineio_path):
try:
ws = websocket.create_connection(
websocket_url + self._get_url_timestamp(), header=headers,
cookie=cookies, enable_multithread=True, **extra_options)
cookie=cookies, enable_multithread=True,
timeout=self.request_timeout, **extra_options)
except (ConnectionError, IOError, websocket.WebSocketException):
if upgrade:
self.logger.warning(
Expand Down
6 changes: 6 additions & 0 deletions tests/asyncio/test_asyncio_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -539,6 +539,7 @@ def test_websocket_connection_failed(self, _time):
c.http.ws_connect.mock.assert_called_once_with(
'ws://foo/engine.io/?transport=websocket&EIO=4&t=123.456',
headers={'Foo': 'Bar'},
timeout=5
)

@mock.patch('engineio.client.time.time', return_value=123.456)
Expand All @@ -553,6 +554,7 @@ def test_websocket_upgrade_failed(self, _time):
c.http.ws_connect.mock.assert_called_once_with(
'ws://foo/engine.io/?transport=websocket&EIO=4&sid=123&t=123.456',
headers={},
timeout=5,
)

def test_websocket_connection_no_open_packet(self):
Expand Down Expand Up @@ -605,6 +607,7 @@ def test_websocket_connection_successful(self, _time):
c.http.ws_connect.mock.assert_called_once_with(
'ws://foo/engine.io/?transport=websocket&EIO=4&t=123.456',
headers={},
timeout=5,
)

@mock.patch('engineio.client.time.time', return_value=123.456)
Expand Down Expand Up @@ -677,6 +680,7 @@ def test_websocket_connection_with_cookies(self, _time):
c.http.ws_connect.mock.assert_called_once_with(
'ws://foo/engine.io/?transport=websocket&EIO=4&t=123.456',
headers={},
timeout=5,
)

@mock.patch('engineio.client.time.time', return_value=123.456)
Expand Down Expand Up @@ -711,6 +715,7 @@ def test_websocket_connection_with_cookie_header(self, _time):
c.http.ws_connect.mock.assert_called_once_with(
'ws://foo/engine.io/?transport=websocket&EIO=4&t=123.456',
headers={},
timeout=5,
)
c.http.cookie_jar.update_cookies.assert_called_once_with(
{'key': 'value', 'key2': 'value2', 'key3': '"value3="'}
Expand Down Expand Up @@ -752,6 +757,7 @@ def test_websocket_connection_with_cookies_and_headers(self, _time):
c.http.ws_connect.mock.assert_called_once_with(
'ws://foo/engine.io/?transport=websocket&EIO=4&t=123.456',
headers={'Foo': 'Bar'},
timeout=5,
)
c.http.cookie_jar.update_cookies.assert_called_once_with(
{'key3': 'value3'}
Expand Down
15 changes: 15 additions & 0 deletions tests/common/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -558,6 +558,7 @@ def test_websocket_connection_failed(self, create_connection, _time):
header={'Foo': 'Bar'},
cookie=None,
enable_multithread=True,
timeout=5
)

@mock.patch('engineio.client.time.time', return_value=123.456)
Expand All @@ -578,6 +579,7 @@ def test_websocket_connection_failed_with_websocket_error(
header={'Foo': 'Bar'},
cookie=None,
enable_multithread=True,
timeout=5
)

@mock.patch('engineio.client.time.time', return_value=123.456)
Expand All @@ -596,6 +598,7 @@ def test_websocket_upgrade_failed(self, create_connection, _time):
header={},
cookie=None,
enable_multithread=True,
timeout=5
)

@mock.patch('engineio.client.websocket.create_connection')
Expand Down Expand Up @@ -644,6 +647,7 @@ def test_websocket_connection_successful(self, create_connection):
'header': {},
'cookie': None,
'enable_multithread': True,
'timeout': 5,
}

@mock.patch('engineio.client.websocket.create_connection')
Expand Down Expand Up @@ -685,6 +689,7 @@ def test_websocket_https_noverify_connection_successful(
'header': {},
'cookie': None,
'enable_multithread': True,
'timeout': 5,
'sslopt': {'cert_reqs': ssl.CERT_NONE},
}

Expand Down Expand Up @@ -722,6 +727,7 @@ def test_websocket_connection_with_cookies(self, create_connection):
'header': {},
'cookie': 'key=value; key2=value2',
'enable_multithread': True,
'timeout': 5,
}

@mock.patch('engineio.client.websocket.create_connection')
Expand Down Expand Up @@ -758,6 +764,7 @@ def test_websocket_connection_with_cookie_header(self, create_connection):
'header': {'Foo': 'bar'},
'cookie': 'key=value',
'enable_multithread': True,
'timeout': 5,
}

@mock.patch('engineio.client.websocket.create_connection')
Expand Down Expand Up @@ -799,6 +806,7 @@ def test_websocket_connection_with_cookies_and_headers(
assert create_connection.call_args[1] == {
'header': {},
'enable_multithread': True,
'timeout': 5,
'cookie': 'key=value; key2=value2; key3=value3',
}

Expand Down Expand Up @@ -831,6 +839,7 @@ def test_websocket_connection_with_auth(self, create_connection):
'header': {'Authorization': 'Basic Zm9vOmJhcg=='},
'cookie': '',
'enable_multithread': True,
'timeout': 5,
}

@mock.patch('engineio.client.websocket.create_connection')
Expand Down Expand Up @@ -863,6 +872,7 @@ def test_websocket_connection_with_cert(self, create_connection):
'header': {},
'cookie': '',
'enable_multithread': True,
'timeout': 5,
}

@mock.patch('engineio.client.websocket.create_connection')
Expand Down Expand Up @@ -895,6 +905,7 @@ def test_websocket_connection_with_cert_and_key(self, create_connection):
'header': {},
'cookie': '',
'enable_multithread': True,
'timeout': 5,
}

@mock.patch('engineio.client.websocket.create_connection')
Expand Down Expand Up @@ -934,6 +945,7 @@ def test_websocket_connection_verify_with_cert_and_key(
'header': {},
'cookie': '',
'enable_multithread': True,
'timeout': 5,
}

@mock.patch('engineio.client.websocket.create_connection')
Expand Down Expand Up @@ -994,6 +1006,7 @@ def test_websocket_connection_with_proxies(self, create_connection):
'header': {},
'cookie': '',
'enable_multithread': True,
'timeout': 5,
}
if results:
expected_results.update({
Expand Down Expand Up @@ -1033,6 +1046,7 @@ def test_websocket_connection_without_verify(self, create_connection):
'header': {},
'cookie': '',
'enable_multithread': True,
'timeout': 5,
}

@mock.patch('engineio.client.websocket.create_connection')
Expand Down Expand Up @@ -1066,6 +1080,7 @@ def test_websocket_connection_with_verify(self, create_connection):
'header': {},
'cookie': '',
'enable_multithread': True,
'timeout': 5,
}

@mock.patch('engineio.client.websocket.create_connection')
Expand Down

0 comments on commit f6df30b

Please sign in to comment.