From a08182a075c9ced33ed3db87e3c10cece502b2e9 Mon Sep 17 00:00:00 2001
From: Matt Mastracci <matthew@mastracci.com>
Date: Tue, 29 Oct 2024 12:55:12 -0600
Subject: [PATCH] Fix exception handling

---
 edb/server/http.py |  3 ++-
 tests/test_http.py | 53 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+), 1 deletion(-)

diff --git a/edb/server/http.py b/edb/server/http.py
index 77d272625198..051b54758c20 100644
--- a/edb/server/http.py
+++ b/edb/server/http.py
@@ -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]]:
@@ -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]
diff --git a/tests/test_http.py b/tests/test_http.py
index f4d026a5deb3..16f967ef6505 100644
--- a/tests/test_http.py
+++ b/tests/test_http.py
@@ -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 = (