test: mmigrate to expect_ where possible (microsoft#428)

Author: pavelfeldman

README.md

@@ -8,7 +8,7 @@ Playwright is a Python library to automate [Chromium](https://www.chromium.org/H
 |   :---   | :---: | :---: | :---:   |
 | Chromium <!-- GEN:chromium-version -->89.0.4344.0<!-- GEN:stop --> | ✅ | ✅ | ✅ |
 | WebKit <!-- GEN:webkit-version -->14.1<!-- GEN:stop --> | ✅ | ✅ | ✅ |
-| Firefox <!-- GEN:firefox-version -->85.0b1<!-- GEN:stop --> | ✅ | ✅ | ✅ |
+| Firefox <!-- GEN:firefox-version -->85.0b5<!-- GEN:stop --> | ✅ | ✅ | ✅ |
 
 Headless execution is supported for all browsers on all platforms.
 

playwright/_impl/_async_base.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import asyncio
-from typing import Any, Callable, Coroutine, Generic, Optional, TypeVar, cast
+from typing import Any, Callable, Generic, TypeVar
 
 from playwright._impl._impl_to_api_mapping import ImplToApiMapping, ImplWrapper
 
@@ -24,22 +24,17 @@
 
 
 class AsyncEventInfo(Generic[T]):
-    def __init__(self, coroutine: Coroutine) -> None:
-        self._value: Optional[T] = None
-        self._future = asyncio.get_event_loop().create_task(coroutine)
-        self._done = False
+    def __init__(self, future: asyncio.Future) -> None:
+        self._future = future
 
     @property
     async def value(self) -> T:
-        if not self._done:
-            self._value = mapping.from_maybe_impl(await self._future)
-            self._done = True
-        return cast(T, self._value)
+        return mapping.from_maybe_impl(await self._future)
 
 
 class AsyncEventContextManager(Generic[T]):
-    def __init__(self, coroutine: Coroutine) -> None:
-        self._event: AsyncEventInfo = AsyncEventInfo(coroutine)
+    def __init__(self, future: asyncio.Future) -> None:
+        self._event: AsyncEventInfo = AsyncEventInfo(future)
 
     async def __aenter__(self) -> AsyncEventInfo[T]:
         return self._event

playwright/_impl/_browser_context.py

@@ -24,7 +24,6 @@
 from playwright._impl._connection import ChannelOwner, from_channel
 from playwright._impl._event_context_manager import EventContextManagerImpl
 from playwright._impl._helper import (
-    PendingWaitEvent,
     RouteHandler,
     RouteHandlerEntry,
     TimeoutSettings,
@@ -55,7 +54,6 @@ def __init__(
         self._pages: List[Page] = []
         self._routes: List[RouteHandlerEntry] = []
         self._bindings: Dict[str, Any] = {}
-        self._pending_wait_for_events: List[PendingWaitEvent] = []
         self._timeout_settings = TimeoutSettings(None)
         self._browser: Optional["Browser"] = None
         self._owner_page: Optional[Page] = None
@@ -201,9 +199,12 @@ async def unroute(
                 "setNetworkInterceptionEnabled", dict(enabled=False)
             )
 
-    async def wait_for_event(
-        self, event: str, predicate: Callable = None, timeout: float = None
-    ) -> Any:
+    def expect_event(
+        self,
+        event: str,
+        predicate: Callable = None,
+        timeout: float = None,
+    ) -> EventContextManagerImpl:
         if timeout is None:
             timeout = self._timeout_settings.timeout()
         wait_helper = WaitHelper(self._loop)
@@ -214,18 +215,14 @@ async def wait_for_event(
             wait_helper.reject_on_event(
                 self, BrowserContext.Events.Close, Error("Context closed")
             )
-        return await wait_helper.wait_for_event(self, event, predicate)
+        wait_helper.wait_for_event(self, event, predicate)
+        return EventContextManagerImpl(wait_helper.result())
 
     def _on_close(self) -> None:
         self._is_closed_or_closing = True
         if self._browser:
             self._browser._contexts.remove(self)
 
-        for pending_event in self._pending_wait_for_events:
-            if pending_event.event == BrowserContext.Events.Close:
-                continue
-            pending_event.reject(False, "Context")
-
         self.emit(BrowserContext.Events.Close)
 
     async def close(self) -> None:
@@ -245,17 +242,16 @@ async def storage_state(self, path: Union[str, Path] = None) -> StorageState:
                 json.dump(result, f)
         return result
 
-    def expect_event(
-        self,
-        event: str,
-        predicate: Callable = None,
-        timeout: float = None,
-    ) -> EventContextManagerImpl:
-        return EventContextManagerImpl(self.wait_for_event(event, predicate, timeout))
+    async def wait_for_event(
+        self, event: str, predicate: Callable = None, timeout: float = None
+    ) -> Any:
+        async with self.expect_event(event, predicate, timeout) as event_info:
+            pass
+        return await event_info.value
 
     def expect_page(
         self,
         predicate: Callable[[Page], bool] = None,
         timeout: float = None,
     ) -> EventContextManagerImpl[Page]:
-        return EventContextManagerImpl(self.wait_for_event("page", predicate, timeout))
+        return self.expect_event(BrowserContext.Events.Page, predicate, timeout)

playwright/_impl/_event_context_manager.py

@@ -13,28 +13,27 @@
 # limitations under the License.
 
 import asyncio
-from typing import Any, Coroutine, Generic, Optional, TypeVar, cast
+from typing import Any, Generic, TypeVar
 
 T = TypeVar("T")
 
 
 class EventInfoImpl(Generic[T]):
-    def __init__(self, coroutine: Coroutine) -> None:
-        self._value: Optional[T] = None
-        self._task = asyncio.get_event_loop().create_task(coroutine)
-        self._done = False
+    def __init__(self, future: asyncio.Future) -> None:
+        self._future = future
 
     @property
     async def value(self) -> T:
-        if not self._done:
-            self._value = await self._task
-            self._done = True
-        return cast(T, self._value)
+        return await self._future
 
 
 class EventContextManagerImpl(Generic[T]):
-    def __init__(self, coroutine: Coroutine) -> None:
-        self._event: EventInfoImpl = EventInfoImpl(coroutine)
+    def __init__(self, future: asyncio.Future) -> None:
+        self._event: EventInfoImpl = EventInfoImpl(future)
+
+    @property
+    def future(self) -> asyncio.Future:
+        return self._event._future
 
     async def __aenter__(self) -> EventInfoImpl[T]:
         return self._event

playwright/_impl/_frame.py

@@ -135,14 +135,14 @@ def _setup_navigation_wait_helper(self, timeout: float = None) -> WaitHelper:
         wait_helper.reject_on_timeout(timeout, f"Timeout {timeout}ms exceeded.")
         return wait_helper
 
-    async def wait_for_navigation(
+    def expect_navigation(
         self,
         url: URLMatch = None,
-        waitUntil: DocumentLoadState = None,
+        wait_until: DocumentLoadState = None,
         timeout: float = None,
-    ) -> Optional[Response]:
-        if not waitUntil:
-            waitUntil = "load"
+    ) -> EventContextManagerImpl[Response]:
+        if not wait_until:
+            wait_until = "load"
 
         if timeout is None:
             timeout = self._page._timeout_settings.navigation_timeout()
@@ -156,23 +156,26 @@ def predicate(event: Any) -> bool:
                 return True
             return not matcher or matcher.matches(event["url"])
 
-        event = await wait_helper.wait_for_event(
+        wait_helper.wait_for_event(
             self._event_emitter,
             "navigated",
             predicate=predicate,
         )
-        if "error" in event:
-            raise Error(event["error"])
-
-        if waitUntil not in self._load_states:
-            t = deadline - monotonic_time()
-            if t > 0:
-                await self.wait_for_load_state(state=waitUntil, timeout=t)
 
-        if "newDocument" in event and "request" in event["newDocument"]:
-            request = from_channel(event["newDocument"]["request"])
-            return await request.response()
-        return None
+        async def continuation() -> Optional[Response]:
+            event = await wait_helper.result()
+            if "error" in event:
+                raise Error(event["error"])
+            if wait_until not in self._load_states:
+                t = deadline - monotonic_time()
+                if t > 0:
+                    await self.wait_for_load_state(state=wait_until, timeout=t)
+            if "newDocument" in event and "request" in event["newDocument"]:
+                request = from_channel(event["newDocument"]["request"])
+                return await request.response()
+            return None
+
+        return EventContextManagerImpl(asyncio.create_task(continuation()))
 
     async def wait_for_load_state(
         self, state: DocumentLoadState = None, timeout: float = None
@@ -184,9 +187,10 @@ async def wait_for_load_state(
         if state in self._load_states:
             return
         wait_helper = self._setup_navigation_wait_helper(timeout)
-        await wait_helper.wait_for_event(
+        wait_helper.wait_for_event(
             self._event_emitter, "loadstate", lambda s: s == state
         )
+        await wait_helper.result()
 
     async def frame_element(self) -> ElementHandle:
         return from_channel(await self._channel.send("frameElement"))
@@ -526,12 +530,14 @@ async def wait_for_function(
     async def title(self) -> str:
         return await self._channel.send("title")
 
-    def expect_navigation(
+    async def wait_for_navigation(
         self,
         url: URLMatch = None,
         waitUntil: DocumentLoadState = None,
         timeout: float = None,
-    ) -> EventContextManagerImpl:
-        return EventContextManagerImpl(
-            self.wait_for_navigation(url, waitUntil, timeout)
-        )
+    ) -> Optional[Response]:
+        async with self.expect_navigation(
+            url, waitUntil, timeout=timeout
+        ) as response_info:
+            pass
+        return await response_info.value

playwright/_impl/_helper.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import asyncio
 import fnmatch
 import math
 import re
@@ -199,25 +198,6 @@ def monotonic_time() -> int:
     return math.floor(time.monotonic() * 1000)
 
 
-class PendingWaitEvent:
-    def __init__(
-        self, event: str, future: asyncio.Future, timeout_future: asyncio.Future
-    ):
-        self.event = event
-        self.future = future
-        self.timeout_future = timeout_future
-
-    def reject(self, is_crash: bool, target: str) -> None:
-        self.timeout_future.cancel()
-        if self.event == "close" and not is_crash:
-            return
-        if self.event == "crash" and is_crash:
-            return
-        self.future.set_exception(
-            Error(f"{target} crashed" if is_crash else f"{target} closed")
-        )
-
-
 class RouteHandlerEntry:
     def __init__(self, matcher: URLMatcher, handler: RouteHandler):
         self.matcher = matcher

playwright/_impl/_network.py

@@ -293,9 +293,12 @@ def __init__(
     def url(self) -> str:
         return self._initializer["url"]
 
-    async def wait_for_event(
-        self, event: str, predicate: Callable = None, timeout: float = None
-    ) -> Any:
+    def expect_event(
+        self,
+        event: str,
+        predicate: Callable = None,
+        timeout: float = None,
+    ) -> EventContextManagerImpl:
         if timeout is None:
             timeout = cast(Any, self._parent)._timeout_settings.timeout()
         wait_helper = WaitHelper(self._loop)
@@ -311,15 +314,15 @@ async def wait_for_event(
                 self, WebSocket.Events.Error, Error("Socket error")
             )
         wait_helper.reject_on_event(self._parent, "close", Error("Page closed"))
-        return await wait_helper.wait_for_event(self, event, predicate)
+        wait_helper.wait_for_event(self, event, predicate)
+        return EventContextManagerImpl(wait_helper.result())
 
-    def expect_event(
-        self,
-        event: str,
-        predicate: Callable = None,
-        timeout: float = None,
-    ) -> EventContextManagerImpl:
-        return EventContextManagerImpl(self.wait_for_event(event, predicate, timeout))
+    async def wait_for_event(
+        self, event: str, predicate: Callable = None, timeout: float = None
+    ) -> Any:
+        async with self.expect_event(event, predicate, timeout) as event_info:
+            pass
+        return await event_info.value
 
     def _on_frame_sent(self, opcode: int, data: str) -> None:
         if opcode == 2: