fix(remote): reject waiters before disconnecting (#3135)

Author: Skn0tt

playwright/_impl/_browser_type.py

@@ -255,15 +255,15 @@ def handle_transport_close(reason: Optional[str]) -> None:
                     for page in context.pages:
                         page._on_close()
                     context._on_close()
-                browser._on_close()
             connection.cleanup(reason)
-            # TODO: Backport https://github.com/microsoft/playwright/commit/d8d5289e8692c9b1265d23ee66988d1ac5122f33
             # Give a chance to any API call promises to reject upon page/context closure.
             # This happens naturally when we receive page.onClose and browser.onClose from the server
             # in separate tasks. However, upon pipe closure we used to dispatch them all synchronously
             # here and promises did not have a chance to reject.
             # The order of rejects vs closure is a part of the API contract and our test runner
             # relies on it to attribute rejections to the right test.
+            if browser:
+                connection._loop.call_soon(browser._on_close)
 
         transport.once("close", handle_transport_close)
 

playwright/_impl/_waiter.py

@@ -17,7 +17,7 @@
 import math
 import uuid
 from asyncio.tasks import Task
-from typing import Any, Callable, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 from pyee import EventEmitter
 
@@ -37,9 +37,7 @@ def __init__(self, channel_owner: ChannelOwner, event: str) -> None:
         self._wait_for_event_info_before(self._wait_id, event)
 
     def _wait_for_event_info_before(self, wait_id: str, event: str) -> None:
-        self._channel.send_no_reply(
-            "__waitInfo__",
-            None,
+        self._send_wait_for_event_info(
             {
                 "waitId": wait_id,
                 "phase": "before",
@@ -49,18 +47,29 @@ def _wait_for_event_info_before(self, wait_id: str, event: str) -> None:
         )
 
     def _wait_for_event_info_after(self, wait_id: str, error: Exception = None) -> None:
-        self._channel._connection.wrap_api_call_sync(
-            lambda: self._channel.send_no_reply(
+        self._send_wait_for_event_info(
+            {
+                "waitId": wait_id,
+                "phase": "after",
+                **({"error": str(error)} if error else {}),
+            },
+            is_internal=True,
+        )
+
+    def _send_wait_for_event_info(
+        self, params: Dict[str, Any], is_internal: bool = False, title: str = None
+    ) -> None:
+        # Wait info is fire-and-forget telemetry and must not affect the waiter.
+        try:
+            self._channel.send_no_reply(
                 "__waitInfo__",
                 None,
-                {
-                    "waitId": wait_id,
-                    "phase": "after",
-                    **({"error": str(error)} if error else {}),
-                },
-            ),
-            True,
-        )
+                params,
+                is_internal=is_internal,
+                title=title,
+            )
+        except Error:
+            pass
 
     def reject_on_event(
         self,
@@ -159,21 +168,14 @@ def result(self) -> asyncio.Future:
 
     def log(self, message: str) -> None:
         self._logs.append(message)
-        try:
-            self._channel._connection.wrap_api_call_sync(
-                lambda: self._channel.send_no_reply(
-                    "__waitInfo__",
-                    None,
-                    {
-                        "waitId": self._wait_id,
-                        "phase": "log",
-                        "message": message,
-                    },
-                ),
-                True,
-            )
-        except Exception:
-            pass
+        self._send_wait_for_event_info(
+            {
+                "waitId": self._wait_id,
+                "phase": "log",
+                "message": message,
+            },
+            is_internal=True,
+        )
 
 
 def throw_on_timeout(timeout: float, exception: Exception) -> asyncio.Task:

tests/async/test_browsertype_connect.py

@@ -180,6 +180,56 @@ async def test_browser_type_connect_should_reject_navigation_when_browser_closes
     assert "has been closed" in exc_info.value.message
 
 
+async def test_browser_type_connect_should_reject_wait_for_event_before_browser_close_finishes(
+    browser_type: BrowserType, launch_server: Callable[[], RemoteServer]
+) -> None:
+    remote_server = launch_server()
+    browser = await browser_type.connect(remote_server.ws_endpoint)
+    page = await browser.new_page()
+    rejected = asyncio.Event()
+
+    async def wait_for_download() -> None:
+        with pytest.raises(Error):
+            await page.wait_for_event("download")
+        rejected.set()
+
+    wait_task = asyncio.create_task(wait_for_download())
+    await asyncio.sleep(0)
+    await browser.close()
+
+    assert rejected.is_set()
+    await wait_task
+
+
+async def test_browser_type_connect_should_reject_wait_for_event_before_disconnected(
+    browser_type: BrowserType, launch_server: Callable[[], RemoteServer]
+) -> None:
+    remote_server = launch_server()
+    browser = await browser_type.connect(remote_server.ws_endpoint)
+    page = await browser.new_page()
+    log = []
+    disconnected = asyncio.Event()
+
+    async def wait_for_download() -> None:
+        with pytest.raises(Error):
+            await page.wait_for_event("download")
+        log.append("rejected")
+
+    def on_disconnected(_: object) -> None:
+        log.append("disconnected")
+        disconnected.set()
+
+    wait_task = asyncio.create_task(wait_for_download())
+    await asyncio.sleep(0)
+    browser.on("disconnected", on_disconnected)
+
+    remote_server.kill()
+    await wait_task
+    await asyncio.wait_for(disconnected.wait(), timeout=5)
+
+    assert log == ["rejected", "disconnected"]
+
+
 async def test_should_not_allow_getting_the_path(
     browser_type: BrowserType, launch_server: Callable[[], RemoteServer], server: Server
 ) -> None: