chore: follow-ups for #1897 connection changes (#1906)

mxschmitt · web-flow · commit 3058517352bc · 2023-05-09T22:11:15.000+02:00
diff --git a/playwright/_impl/_browser_type.py b/playwright/_impl/_browser_type.py
@@ -219,7 +219,7 @@ async def connect(
 
         timeout_future = throw_on_timeout(timeout, Error("Connection timed out"))
         done, pending = await asyncio.wait(
-            {playwright_future, timeout_future},
+            {transport.on_error_future, playwright_future, timeout_future},
             return_when=asyncio.FIRST_COMPLETED,
         )
         if not playwright_future.done():
@@ -235,13 +235,13 @@ async def connect(
         self._did_launch_browser(browser)
         browser._should_close_connection_on_close = True
 
-        def handle_transport_close(transport_exception: str) -> None:
+        def handle_transport_close() -> None:
             for context in browser.contexts:
                 for page in context.pages:
                     page._on_close()
                 context._on_close()
             browser._on_close()
-            connection.cleanup(transport_exception or BROWSER_CLOSED_ERROR)
+            connection.cleanup(BROWSER_CLOSED_ERROR)
 
         transport.once("close", handle_transport_close)
 
diff --git a/playwright/_impl/_connection.py b/playwright/_impl/_connection.py
@@ -290,8 +290,6 @@ def cleanup(self, error_message: str = None) -> None:
             ws_connection._transport.dispose()
         for callback in self._callbacks.values():
             callback.future.set_exception(Error(error_message))
-            # Prevent 'Task exception was never retrieved'
-            callback.future.exception()
         self._callbacks.clear()
         self.emit("close")
 
diff --git a/playwright/_impl/_json_pipe.py b/playwright/_impl/_json_pipe.py
@@ -13,12 +13,12 @@
 # limitations under the License.
 
 import asyncio
-from typing import Dict, Optional, cast
+from typing import Dict, cast
 
 from pyee.asyncio import AsyncIOEventEmitter
 
 from playwright._impl._connection import Channel
-from playwright._impl._helper import ParsedMessagePayload
+from playwright._impl._helper import Error, ParsedMessagePayload
 from playwright._impl._transport import Transport
 
 
@@ -48,17 +48,13 @@ async def wait_until_stopped(self) -> None:
     async def connect(self) -> None:
         self._stopped_future: asyncio.Future = asyncio.Future()
 
-        close_error: Optional[str] = None
-
         def handle_message(message: Dict) -> None:
-            try:
-                self.on_message(cast(ParsedMessagePayload, message))
-            except Exception as e:
-                nonlocal close_error
-                close_error = str(e)
+            if self._stop_requested:
+                return
+            self.on_message(cast(ParsedMessagePayload, message))
 
         def handle_closed() -> None:
-            self.emit("close", close_error)
+            self.emit("close")
             self._stopped_future.set_result(None)
 
         self._pipe_channel.on(
@@ -74,4 +70,6 @@ async def run(self) -> None:
         await self._stopped_future
 
     def send(self, message: Dict) -> None:
+        if self._stop_requested:
+            raise Error("Playwright connection closed")
         self._pipe_channel.send_no_reply("send", {"message": message})
