Skip to content

Commit

Permalink
Fix exception handling
Browse files Browse the repository at this point in the history
  • Loading branch information
mmastrac committed Oct 29, 2024
1 parent 57db287 commit a08182a
Show file tree
Hide file tree
Showing 2 changed files with 55 additions and 1 deletion.
3 changes: 2 additions & 1 deletion edb/server/http.py
Original file line number Diff line number Diff line change
Expand Up @@ -119,6 +119,7 @@ def _safe_ack(self, id):

def _update_limit(self, limit: int):
if self._client is not None and limit != self._limit:
self._limit = limit
self._client._update_limit(limit)

def _process_headers(self, headers: HeaderType) -> list[tuple[str, str]]:
Expand Down Expand Up @@ -320,7 +321,7 @@ def _process_message(self, msg):
msg_type, id, data = msg
if msg_type == 0: # Error
if id in self._requests:
self._requests[id].set_exception(Exception(data[0]))
self._requests[id].set_exception(Exception(data))
if id in self._streaming:
self._streaming[id].put_nowait(None)
del self._streaming[id]
Expand Down
53 changes: 53 additions & 0 deletions tests/test_http.py
Original file line number Diff line number Diff line change
Expand Up @@ -112,6 +112,59 @@ async def test_post_with_headers(self):
)
self.assertEqual(result.headers["X-Test"], "test!")

async def test_bad_url(self):
with http.HttpClient(100) as client:
with self.assertRaisesRegex(Exception, "Scheme"):
await client.get("httpx://uh-oh")

async def test_immediate_connection_drop(self):
"""Test handling of a connection that is dropped immediately by the server"""

async def mock_drop_server(
_reader: asyncio.StreamReader, writer: asyncio.StreamWriter
):
# Close connection immediately without sending any response
writer.close()
await writer.wait_closed()

server = await asyncio.start_server(mock_drop_server, 'localhost', 0)
addr = server.sockets[0].getsockname()
url = f'http://{addr[0]}:{addr[1]}/drop'

try:
with http.HttpClient(100) as client:
with self.assertRaisesRegex(
Exception, "Connection reset by peer"
):
await client.get(url)
finally:
server.close()
await server.wait_closed()

async def test_immediate_connection_drop_streaming(self):
"""Test handling of a connection that is dropped immediately by the server"""

async def mock_drop_server(
_reader: asyncio.StreamReader, writer: asyncio.StreamWriter
):
# Close connection immediately without sending any response
writer.close()
await writer.wait_closed()

server = await asyncio.start_server(mock_drop_server, 'localhost', 0)
addr = server.sockets[0].getsockname()
url = f'http://{addr[0]}:{addr[1]}/drop'

try:
with http.HttpClient(100) as client:
with self.assertRaisesRegex(
Exception, "Connection reset by peer"
):
await client.stream_sse(url)
finally:
server.close()
await server.wait_closed()

async def test_streaming_get_with_no_sse(self):
with http.HttpClient(100) as client:
example_request = (
Expand Down

0 comments on commit a08182a

Please sign in to comment.