enhancements

pytest_playwright/pytest_playwright.py

@@ -18,7 +18,20 @@
 import sys
 import warnings
 from pathlib import Path
-from typing import Any, Callable, Dict, Generator, List, Optional
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Generator,
+    List,
+    Literal,
+    Optional,
+    Protocol,
+    Sequence,
+    Union,
+    Pattern,
+    cast,
+)
 
 import pytest
 from playwright.sync_api import (
@@ -29,11 +42,15 @@
     Page,
     Playwright,
     sync_playwright,
+    ProxySettings,
+    StorageState,
+    HttpCredentials,
+    Geolocation,
+    ViewportSize,
 )
 from slugify import slugify
 import tempfile
 
-
 artifacts_folder = tempfile.TemporaryDirectory(prefix="playwright-pytest-")
 
 
@@ -243,19 +260,62 @@ def browser(launch_browser: Callable[[], Browser]) -> Generator[Browser, None, N
         pass
 
 
+class CreateContextCallback(Protocol):
+    def __call__(
+        self,
+        viewport: Optional[ViewportSize] = None,
+        screen: Optional[ViewportSize] = None,
+        no_viewport: Optional[bool] = None,
+        ignore_https_errors: Optional[bool] = None,
+        java_script_enabled: Optional[bool] = None,
+        bypass_csp: Optional[bool] = None,
+        user_agent: Optional[str] = None,
+        locale: Optional[str] = None,
+        timezone_id: Optional[str] = None,
+        geolocation: Optional[Geolocation] = None,
+        permissions: Optional[Sequence[str]] = None,
+        extra_http_headers: Optional[Dict[str, str]] = None,
+        offline: Optional[bool] = None,
+        http_credentials: Optional[HttpCredentials] = None,
+        device_scale_factor: Optional[float] = None,
+        is_mobile: Optional[bool] = None,
+        has_touch: Optional[bool] = None,
+        color_scheme: Optional[
+            Literal["dark", "light", "no-preference", "null"]
+        ] = None,
+        reduced_motion: Optional[Literal["no-preference", "null", "reduce"]] = None,
+        forced_colors: Optional[Literal["active", "none", "null"]] = None,
+        accept_downloads: Optional[bool] = None,
+        default_browser_type: Optional[str] = None,
+        proxy: Optional[ProxySettings] = None,
+        record_har_path: Optional[Union[str, Path]] = None,
+        record_har_omit_content: Optional[bool] = None,
+        record_video_dir: Optional[Union[str, Path]] = None,
+        record_video_size: Optional[ViewportSize] = None,
+        storage_state: Optional[Union[StorageState, str, Path]] = None,
+        base_url: Optional[str] = None,
+        strict_selectors: Optional[bool] = None,
+        service_workers: Optional[Literal["allow", "block"]] = None,
+        record_har_url_filter: Optional[Union[str, Pattern[str]]] = None,
+        record_har_mode: Optional[Literal["full", "minimal"]] = None,
+        record_har_content: Optional[Literal["attach", "embed", "omit"]] = None,
+    ) -> BrowserContext:
+        ...
+
+
 @pytest.fixture
 def new_context(
     browser: Browser,
     browser_context_args: Dict,
     _artifacts_recorder: "ArtifactsRecorder",
     request: pytest.FixtureRequest,
-) -> Callable[..., BrowserContext]:
+) -> CreateContextCallback:
     browser_context_args = browser_context_args.copy()
     context_args_marker = next(request.node.iter_markers("browser_context_args"), None)
     additional_context_args = context_args_marker.kwargs if context_args_marker else {}
     browser_context_args.update(additional_context_args)
 
-    def _new_context(**kwargs: Dict) -> BrowserContext:
+    def _new_context(**kwargs: Any) -> BrowserContext:
         context = browser.new_context(**browser_context_args, **kwargs)
         original_close = context.close
 
@@ -267,21 +327,17 @@ def close_wrapper(*args: Any, **kwargs: Any) -> None:
         _artifacts_recorder.on_did_create_browser_context(context)
         return context
 
-    return _new_context
+    return cast(CreateContextCallback, _new_context)
 
 
 @pytest.fixture
-def context(
-    new_context: Callable[..., BrowserContext]
-) -> Generator[BrowserContext, None, None]:
-    context = new_context()
-    yield context
+def context(new_context: CreateContextCallback) -> BrowserContext:
+    return new_context()
 
 
 @pytest.fixture
-def page(context: BrowserContext) -> Generator[Page, None, None]:
-    page = context.new_page()
-    yield page
+def page(context: BrowserContext) -> Page:
+    return context.new_page()
 
 
 @pytest.fixture(scope="session")
@@ -396,8 +452,8 @@ def __init__(
         self._pytestconfig = pytestconfig
         self._playwright = playwright
 
-        self._contexts: BrowserContext = []
-        self._pages: List[Page] = []
+        self._open_contexts: BrowserContext = []
+        self._all_pages: List[Page] = []
         self._traces: List[str] = []
         self._tracing_option = pytestconfig.getoption("--tracing")
         self._capture_trace = self._tracing_option in ["on", "retain-on-failure"]
@@ -408,7 +464,7 @@ def did_finish_test(self, failed: bool) -> None:
             failed and screenshot_option == "only-on-failure"
         )
         if capture_screenshot:
-            for index, page in enumerate(self._pages):
+            for index, page in enumerate(self._all_pages):
                 human_readable_status = "failed" if failed else "finished"
                 screenshot_path = _build_artifact_test_folder(
                     self._pytestconfig,
@@ -427,8 +483,8 @@ def did_finish_test(self, failed: bool) -> None:
                     pass
 
         # Close contexts which were not closed during the test (this will trigger Trace and Video generation)
-        while len(self._contexts) > 0:
-            self._contexts[0].close()
+        while len(self._open_contexts) > 0:
+            self._open_contexts[0].close()
 
         if self._tracing_option == "on" or (
             failed and self._tracing_option == "retain-on-failure"
@@ -452,14 +508,14 @@ def did_finish_test(self, failed: bool) -> None:
             failed and video_option == "retain-on-failure"
         )
         if preserve_video:
-            for index, page in enumerate(self._pages):
+            for index, page in enumerate(self._all_pages):
                 video = page.video
                 if not video:
                     continue
                 try:
                     video_file_name = (
                         "video.webm"
-                        if len(self._pages) == 1
+                        if len(self._all_pages) == 1
                         else f"video-{index+1}.webm"
                     )
                     video.save_as(
@@ -472,8 +528,8 @@ def did_finish_test(self, failed: bool) -> None:
                     pass
 
     def on_did_create_browser_context(self, context: BrowserContext) -> None:
-        self._contexts.append(context)
-        context.on("page", lambda page: self._pages.append(page))
+        self._open_contexts.append(context)
+        context.on("page", lambda page: self._all_pages.append(page))
         if self._request and self._capture_trace:
             context.tracing.start(
                 title=slugify(self._request.node.nodeid),
@@ -483,8 +539,8 @@ def on_did_create_browser_context(self, context: BrowserContext) -> None:
             )
 
     def on_will_close_browser_context(self, context: BrowserContext) -> None:
-        if context in self._contexts:
-            self._contexts.remove(context)
+        if context in self._open_contexts:
+            self._open_contexts.remove(context)
         if self._capture_trace:
             trace_path = Path(artifacts_folder.name) / create_guid()
             context.tracing.stop(path=trace_path)