feat(wait helper): add wait helper for client-side waitFor* methods (microsoft#74)

Author: pavelfeldman

playwright/browser_context.py

@@ -29,7 +29,8 @@
     URLMatcher,
 )
 from playwright.network import Request, Route, serialize_headers
-from playwright.page import BindingCall, Page, wait_for_event
+from playwright.page import BindingCall, Page
+from playwright.wait_helper import WaitHelper
 
 if TYPE_CHECKING:  # pragma: no cover
     from playwright.browser import Browser
@@ -185,9 +186,17 @@ async def unroute(
     async def waitForEvent(
         self, event: str, predicate: Callable[[Any], bool] = None, timeout: int = None
     ) -> Any:
-        return await wait_for_event(
-            self, self._timeout_settings, event, predicate=predicate, timeout=timeout
+        if timeout is None:
+            timeout = self._timeout_settings.timeout()
+        wait_helper = WaitHelper()
+        wait_helper.reject_on_timeout(
+            timeout, f'Timeout while waiting for event "${event}"'
         )
+        if event != BrowserContext.Events.Close:
+            wait_helper.reject_on_event(
+                self, BrowserContext.Events.Close, Error("Context closed")
+            )
+        return await wait_helper.wait_for_event(self, event, predicate)
 
     def _on_close(self) -> None:
         self._is_closed_or_closing = True

playwright/frame.py

@@ -32,14 +32,19 @@
 )
 from playwright.helper import (
     DocumentLoadState,
+    Error,
     FilePayload,
+    FrameNavigatedEvent,
     KeyboardModifier,
     MouseButton,
+    URLMatch,
+    URLMatcher,
     is_function_body,
     locals_to_params,
 )
 from playwright.js_handle import JSHandle, parse_result, serialize_argument
 from playwright.network import Response
+from playwright.wait_helper import WaitHelper
 
 if sys.version_info >= (3, 8):  # pragma: no cover
     from typing import Literal
@@ -68,27 +73,30 @@ def __init__(self, scope: ConnectionScope, guid: str, initializer: Dict) -> None
             lambda params: self._on_load_state(params.get("add"), params.get("remove")),
         )
         self._channel.on(
-            "navigated",
-            lambda params: self._on_frame_navigated(params["url"], params["name"]),
+            "navigated", lambda params: self._on_frame_navigated(params),
         )
 
-    def _on_load_state(self, add: str = None, remove: str = None) -> None:
+    def _on_load_state(
+        self, add: DocumentLoadState = None, remove: DocumentLoadState = None
+    ) -> None:
         if add:
             self._load_states.add(add)
             self._event_emitter.emit("loadstate", add)
         elif remove and remove in self._load_states:
             self._load_states.remove(remove)
 
-    def _on_frame_navigated(self, url: str, name: str) -> None:
-        self._url = url
-        self._name = name
-        self._page.emit("framenavigated", self)
+    def _on_frame_navigated(self, event: FrameNavigatedEvent) -> None:
+        self._url = event["url"]
+        self._name = event["name"]
+        self._event_emitter.emit("navigated", event)
+        if "error" not in event and self._page:
+            self._page.emit("framenavigated", self)
 
     async def goto(
         self,
         url: str,
         timeout: int = None,
-        waitUntil: DocumentLoadState = None,
+        waitUntil: DocumentLoadState = "load",
         referer: str = None,
     ) -> Optional[Response]:
         return cast(
@@ -98,33 +106,54 @@ async def goto(
             ),
         )
 
+    def _setup_navigation_wait_helper(self, timeout: int = None) -> WaitHelper:
+        wait_helper = WaitHelper()
+        wait_helper.reject_on_event(
+            self._page, "close", Error("Navigation failed because page was closed!")
+        )
+        wait_helper.reject_on_event(
+            self._page, "crash", Error("Navigation failed because page crashed!")
+        )
+        wait_helper.reject_on_event(
+            self._page,
+            "framedetached",
+            Error("Navigating frame was detached!"),
+            lambda frame: frame == self,
+        )
+        if timeout is None:
+            timeout = self._page._timeout_settings.navigation_timeout()
+        wait_helper.reject_on_timeout(timeout, f"Timeout {timeout}ms exceeded.")
+        return wait_helper
+
     async def waitForNavigation(
         self,
         timeout: int = None,
-        waitUntil: DocumentLoadState = None,
-        url: str = None,  # TODO: add url, callback
+        waitUntil: DocumentLoadState = "load",
+        url: URLMatch = None,
     ) -> Optional[Response]:
-        return cast(
-            Optional[Response],
-            from_nullable_channel(
-                await self._channel.send(
-                    "waitForNavigation", locals_to_params(locals())
-                )
-            ),
+        wait_helper = self._setup_navigation_wait_helper(timeout)
+        matcher = URLMatcher(url) if url else None
+        event = await wait_helper.wait_for_event(
+            self._event_emitter,
+            "navigated",
+            lambda event: not matcher or matcher.matches(event["url"]),
         )
-
-    async def waitForLoadState(self, state: str = "load", timeout: int = None) -> None:
+        if "newDocument" in event and "request" in event["newDocument"]:
+            request = from_channel(event["newDocument"]["request"])
+            return await request.response()
+        if "error" in event:
+            raise Error(event["error"])
+        return None
+
+    async def waitForLoadState(
+        self, state: DocumentLoadState = "load", timeout: int = None
+    ) -> None:
         if state in self._load_states:
             return
-        future = self._scope._loop.create_future()
-
-        def loadstate(s: str) -> None:
-            if state == s:
-                future.set_result(None)
-                self._event_emitter.remove_listener("loadstate", loadstate)
-
-        self._event_emitter.on("loadstate", loadstate)
-        await future
+        wait_helper = self._setup_navigation_wait_helper(timeout)
+        await wait_helper.wait_for_event(
+            self._event_emitter, "loadstate", lambda s: s == state
+        )
 
     async def frameElement(self) -> ElementHandle:
         return from_channel(await self._channel.send("frameElement"))

playwright/helper.py

@@ -106,6 +106,17 @@ class ParsedMessagePayload(TypedDict, total=False):
     error: ErrorPayload
 
 
+class Document(TypedDict):
+    request: Optional[Any]
+
+
+class FrameNavigatedEvent(TypedDict):
+    url: str
+    name: str
+    newDocument: Optional[Document]
+    error: Optional[str]
+
+
 class URLMatcher:
     def __init__(self, match: URLMatch) -> None:
         self._callback: Optional[Callable[[str], bool]] = None

playwright/network.py

@@ -59,7 +59,6 @@ def postData(self) -> Optional[str]:
     def headers(self) -> Dict[str, str]:
         return parse_headers(self._initializer["headers"])
 
-    @property
     async def response(self) -> Optional["Response"]:
         return from_nullable_channel(await self._channel.send("response"))
 

playwright/page.py

@@ -40,7 +40,6 @@
     PendingWaitEvent,
     RouteHandler,
     RouteHandlerEntry,
-    TimeoutError,
     TimeoutSettings,
     URLMatch,
     URLMatcher,
@@ -53,6 +52,7 @@
 from playwright.input import Keyboard, Mouse
 from playwright.js_handle import JSHandle, serialize_argument
 from playwright.network import Request, Response, Route, serialize_headers
+from playwright.wait_helper import WaitHelper
 from playwright.worker import Worker
 
 if sys.version_info >= (3, 8):  # pragma: no cover
@@ -103,6 +103,7 @@ def __init__(self, scope: ConnectionScope, guid: str, initializer: Dict) -> None
         self._pending_wait_for_events: List[PendingWaitEvent] = list()
         self._routes: List[RouteHandlerEntry] = list()
         self._owned_context: Optional["BrowserContext"] = None
+        self._timeout_settings = TimeoutSettings(None)
 
         self._channel.on(
             "bindingCall",
@@ -370,25 +371,27 @@ async def goto(
         self,
         url: str,
         timeout: int = None,
-        waitUntil: DocumentLoadState = None,
+        waitUntil: DocumentLoadState = "load",
         referer: str = None,
     ) -> Optional[Response]:
         return await self._main_frame.goto(**locals_to_params(locals()))
 
     async def reload(
-        self, timeout: int = None, waitUntil: DocumentLoadState = None,
+        self, timeout: int = None, waitUntil: DocumentLoadState = "load",
     ) -> Optional[Response]:
         return from_nullable_channel(
             await self._channel.send("reload", locals_to_params(locals()))
         )
 
-    async def waitForLoadState(self, state: str = "load", timeout: int = None) -> None:
+    async def waitForLoadState(
+        self, state: DocumentLoadState = "load", timeout: int = None
+    ) -> None:
         return await self._main_frame.waitForLoadState(**locals_to_params(locals()))
 
     async def waitForNavigation(
         self,
         timeout: int = None,
-        waitUntil: DocumentLoadState = None,
+        waitUntil: DocumentLoadState = "load",
         url: str = None,  # TODO: add url, callback
     ) -> Optional[Response]:
         return await self._main_frame.waitForNavigation(**locals_to_params(locals()))
@@ -440,9 +443,17 @@ def my_predicate(response: Response) -> bool:
     async def waitForEvent(
         self, event: str, predicate: Callable[[Any], bool] = None, timeout: int = None
     ) -> Any:
-        return await wait_for_event(
-            self, self._timeout_settings, event, predicate=predicate, timeout=timeout
+        if timeout is None:
+            timeout = self._timeout_settings.timeout()
+        wait_helper = WaitHelper()
+        wait_helper.reject_on_timeout(
+            timeout, f'Timeout while waiting for event "${event}"'
         )
+        if event != Page.Events.Crash:
+            wait_helper.reject_on_event(self, Page.Events.Crash, Error("Page crashed"))
+        if event != Page.Events.Close:
+            wait_helper.reject_on_event(self, Page.Events.Close, Error("Page closed"))
+        return await wait_helper.wait_for_event(self, event, predicate)
 
     async def goBack(
         self, timeout: int = None, waitUntil: DocumentLoadState = None,
@@ -480,20 +491,19 @@ async def addInitScript(self, source: str = None, path: str = None) -> None:
             raise Error("Either path or source parameter must be specified")
         await self._channel.send("addInitScript", dict(source=source))
 
-    async def route(self, match: URLMatch, handler: RouteHandler) -> None:
-        self._routes.append(RouteHandlerEntry(URLMatcher(match), handler))
+    async def route(self, url: URLMatch, handler: RouteHandler) -> None:
+        self._routes.append(RouteHandlerEntry(URLMatcher(url), handler))
         if len(self._routes) == 1:
             await self._channel.send(
                 "setNetworkInterceptionEnabled", dict(enabled=True)
             )
 
     async def unroute(
-        self, match: URLMatch, handler: Optional[RouteHandler] = None
+        self, url: URLMatch, handler: Optional[RouteHandler] = None
     ) -> None:
         self._routes = list(
             filter(
-                lambda r: r.matcher.match != match
-                or (handler and r.handler != handler),
+                lambda r: r.matcher.match != url or (handler and r.handler != handler),
                 self._routes,
             )
         )
@@ -713,39 +723,3 @@ async def call(self, func: FunctionWithSource) -> None:
             asyncio.ensure_future(
                 self._channel.send("reject", dict(error=serialize_error(e, tb)))
             )
-
-
-async def wait_for_event(
-    target: Union[Page, "BrowserContext"],
-    timeout_settings: TimeoutSettings,
-    event: str,
-    predicate: Callable[[Any], bool] = None,
-    timeout: int = None,
-) -> Any:
-    if timeout is None:
-        timeout = timeout_settings.timeout()
-    if timeout == 0:
-        timeout = 3600 * 24 * 7 * 30 * 365
-    timeout_future: asyncio.Future = asyncio.ensure_future(
-        asyncio.sleep(timeout / 1000)
-    )
-
-    future = target._scope._loop.create_future()
-
-    def listener(e: Any = None) -> None:
-        if not predicate or predicate(e):
-            future.set_result(e)
-
-    target.on(event, listener)
-
-    pending_event = PendingWaitEvent(event, future, timeout_future)
-    target._pending_wait_for_events.append(pending_event)
-    done, _ = await asyncio.wait(
-        {timeout_future, future}, return_when=asyncio.FIRST_COMPLETED
-    )
-    target.remove_listener(event, listener)
-    target._pending_wait_for_events.remove(pending_event)
-    if future in done:
-        timeout_future.cancel()
-        return future.result()
-    raise TimeoutError("Timeout exceeded")