wip

Author: pavelfeldman

playwright/_impl/_greenlets.py

@@ -68,9 +68,13 @@ def wait_for_future(
 ) -> None:
     __tracebackhide__ = True
     g_self = greenlet.getcurrent()
+    waiting = True
 
     def resume(_: "asyncio.Future") -> None:
-        g_self.switch()
+        # asyncio runs the callback in a later loop iteration, also when the future
+        # is done already. By then the wait may be over and this greenlet elsewhere.
+        if waiting:
+            g_self.switch()
 
     # The wait owns the callback that resumes this greenlet, so that it can
     # remove it again when the wait is abandoned. A callback left behind would
@@ -92,6 +96,7 @@ def resume(_: "asyncio.Future") -> None:
         future.cancel()
         raise
     finally:
+        waiting = False
         # The loop is only running for as long as the dispatcher fiber is alive.
         # Also after a signal handler interrupted it: run_forever() cleared this
         # on its way out, but the dispatcher is going to run the loop again.

tests/async/library/test_browsertype_connect.py

@@ -493,9 +493,8 @@ async def test_BT_046_should_send_extra_headers_with_connect_request(
     kind: str,
     connect: Callable[..., Awaitable[Browser]],
     server: Server,
-    server_web_sockets: List[WebSocketProtocol],
 ) -> None:
-    server.once_web_socket_connection(server_web_sockets.append)
+    request_future = server.wait_for_web_socket_connection_request()
     with suppress(Error):
         await connect(
             f"ws://localhost:{server.PORT}/ws",
@@ -505,9 +504,9 @@ async def test_BT_046_should_send_extra_headers_with_connect_request(
             },
             timeout=100,
         )
-    request = server_web_sockets[0].http_headers
-    assert request["user-agent"] == "Playwright"
-    assert request["foo"] == "bar"
+    request = await request_future
+    assert request.headers["user-agent"] == "Playwright"
+    assert request.headers["foo"] == "bar"
 
 
 @pytest.mark.parametrize("kind", KINDS)

tests/sync/legacy/test_sync.py

@@ -304,26 +304,39 @@ def test_expect_response_should_not_hang_when_predicate_throws(page: Page) -> No
             raise Exception("Oops!")
 
 
-def test_expect_response_should_not_switch_to_stale_greenlet_when_cancelled(
+def test_expect_event_should_not_switch_to_stale_greenlet(
     page: Page, server: Server
 ) -> None:
     # Mimics frameworks like pytest-twisted that run the test body in a separate
     # greenlet which stays alive (suspended) after handing control back.
     main = greenlet.getcurrent()
-    resumed: List[bool] = []
 
-    def body() -> None:
+    def expect_not_resumed(action: Callable[[], None]) -> None:
+        resumed: List[bool] = []
+
+        def body() -> None:
+            action()
+            main.switch()
+            resumed.append(True)
+            main.switch()
+
+        greenlet(body).switch()
+        # Pumps the dispatcher, which runs the done callbacks of the futures.
+        page.goto(server.EMPTY_PAGE)
+        assert resumed == []
+
+    def cancel_wait() -> None:
         with pytest.raises(Exception, match="Oops!"):
             with page.expect_response("**/*"):
                 raise Exception("Oops!")
-        main.switch()
-        resumed.append(True)
-        main.switch()
 
-    greenlet(body).switch()
-    # Pumps the dispatcher, which runs the cancelled future's done callbacks.
-    page.goto(server.EMPTY_PAGE)
-    assert resumed == []
+    def wait_for_done_future() -> None:
+        with page.expect_console_message() as message_info:
+            page.evaluate("console.log('hello')")
+        assert message_info.value.text == "hello"
+
+    expect_not_resumed(cancel_wait)
+    expect_not_resumed(wait_for_done_future)
 
 
 def test_expect_response_should_use_context_timeout(

tests/sync/library/test_browsertype_connect.py

@@ -506,11 +506,11 @@ def test_BT_045_should_timeout_in_connect_while_connecting(
 @pytest.mark.parametrize("kind", KINDS)
 def test_BT_046_should_send_extra_headers_with_connect_request(
     kind: str,
+    playwright: Playwright,
     connect: Callable[..., Browser],
     server: Server,
-    server_web_sockets: List[WebSocketProtocol],
 ) -> None:
-    server.once_web_socket_connection(server_web_sockets.append)
+    request_future = server.wait_for_web_socket_connection_request()
     with suppress(Error):
         connect(
             f"ws://localhost:{server.PORT}/ws",
@@ -520,9 +520,10 @@ def test_BT_046_should_send_extra_headers_with_connect_request(
             },
             timeout=100,
         )
-    request = server_web_sockets[0].http_headers
-    assert request["user-agent"] == "Playwright"
-    assert request["foo"] == "bar"
+    wait_for_future(playwright._loop, playwright._dispatcher_fiber, request_future)
+    request = request_future.result()
+    assert request.headers["user-agent"] == "Playwright"
+    assert request.headers["foo"] == "bar"
 
 
 @pytest.mark.parametrize("kind", KINDS)