feat: support artifact collection with multiple contexts

pytest_playwright/pytest_playwright.py

@@ -17,7 +17,21 @@
 import os
 import sys
 import warnings
-from typing import Any, Callable, Dict, Generator, List, Optional
+from pathlib import Path
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Generator,
+    List,
+    Literal,
+    Optional,
+    Protocol,
+    Sequence,
+    Union,
+    Pattern,
+    cast,
+)
 
 import pytest
 from playwright.sync_api import (
@@ -28,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-")
 
 
@@ -190,6 +208,20 @@ def browser_context_args(
     return context_args
 
 
+@pytest.fixture()
+def _artifacts_recorder(
+    request: pytest.FixtureRequest,
+    playwright: Playwright,
+    pytestconfig: Any,
+) -> Generator["ArtifactsRecorder", None, None]:
+    artifacts_recorder = ArtifactsRecorder(pytestconfig, request, playwright)
+    yield artifacts_recorder
+    # If request.node is missing rep_call, then some error happened during execution
+    # that prevented teardown, but should still be counted as a failure
+    failed = request.node.rep_call.failed if hasattr(request.node, "rep_call") else True
+    artifacts_recorder.did_finish_test(failed)
+
+
 @pytest.fixture(scope="session")
 def playwright() -> Generator[Playwright, None, None]:
     pw = sync_playwright().start()
@@ -228,93 +260,89 @@ 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 context(
+def new_context(
     browser: Browser,
     browser_context_args: Dict,
-    pytestconfig: Any,
+    _artifacts_recorder: "ArtifactsRecorder",
     request: pytest.FixtureRequest,
-) -> Generator[BrowserContext, None, None]:
-    pages: List[Page] = []
-
+) -> Generator[CreateContextCallback, None, None]:
     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)
+    contexts: List[BrowserContext] = []
 
-    context = browser.new_context(**browser_context_args)
-    context.on("page", lambda page: pages.append(page))
-
-    tracing_option = pytestconfig.getoption("--tracing")
-    capture_trace = tracing_option in ["on", "retain-on-failure"]
-    if capture_trace:
-        context.tracing.start(
-            title=slugify(request.node.nodeid),
-            screenshots=True,
-            snapshots=True,
-            sources=True,
-        )
+    def _new_context(**kwargs: Any) -> BrowserContext:
+        context = browser.new_context(**browser_context_args, **kwargs)
+        original_close = context.close
 
-    yield context
+        def _close_wrapper(*args: Any, **kwargs: Any) -> None:
+            contexts.remove(context)
+            _artifacts_recorder.on_will_close_browser_context(context)
+            original_close(*args, **kwargs)
 
-    # If request.node is missing rep_call, then some error happened during execution
-    # that prevented teardown, but should still be counted as a failure
-    failed = request.node.rep_call.failed if hasattr(request.node, "rep_call") else True
-
-    if capture_trace:
-        retain_trace = tracing_option == "on" or (
-            failed and tracing_option == "retain-on-failure"
-        )
-        if retain_trace:
-            trace_path = _build_artifact_test_folder(pytestconfig, request, "trace.zip")
-            context.tracing.stop(path=trace_path)
-        else:
-            context.tracing.stop()
+        context.close = _close_wrapper
+        contexts.append(context)
+        _artifacts_recorder.on_did_create_browser_context(context)
+        return context
 
-    screenshot_option = pytestconfig.getoption("--screenshot")
-    capture_screenshot = screenshot_option == "on" or (
-        failed and screenshot_option == "only-on-failure"
-    )
-    if capture_screenshot:
-        for index, page in enumerate(pages):
-            human_readable_status = "failed" if failed else "finished"
-            screenshot_path = _build_artifact_test_folder(
-                pytestconfig, request, f"test-{human_readable_status}-{index+1}.png"
-            )
-            try:
-                page.screenshot(
-                    timeout=5000,
-                    path=screenshot_path,
-                    full_page=pytestconfig.getoption("--full-page-screenshot"),
-                )
-            except Error:
-                pass
+    yield cast(CreateContextCallback, _new_context)
+    for context in contexts.copy():
+        context.close()
 
-    context.close()
 
-    video_option = pytestconfig.getoption("--video")
-    preserve_video = video_option == "on" or (
-        failed and video_option == "retain-on-failure"
-    )
-    if preserve_video:
-        for i, page in enumerate(pages):
-            video = page.video
-            if not video:
-                continue
-            try:
-                video_name = "video.webm" if len(pages) == 1 else f"video-{i+1}.webm"
-                video.save_as(
-                    path=_build_artifact_test_folder(pytestconfig, request, video_name)
-                )
-            except Error:
-                # Silent catch empty videos.
-                pass
+@pytest.fixture
+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")
@@ -419,3 +447,114 @@ def pytest_addoption(parser: Any) -> None:
         default=False,
         help="Whether to take a full page screenshot",
     )
+
+
+class ArtifactsRecorder:
+    def __init__(
+        self, pytestconfig: Any, request: pytest.FixtureRequest, playwright: Playwright
+    ) -> None:
+        self._request = request
+        self._pytestconfig = pytestconfig
+        self._playwright = playwright
+
+        self._all_pages: List[Page] = []
+        self._screenshots: List[str] = []
+        self._traces: List[str] = []
+        self._tracing_option = pytestconfig.getoption("--tracing")
+        self._capture_trace = self._tracing_option in ["on", "retain-on-failure"]
+
+    def did_finish_test(self, failed: bool) -> None:
+        screenshot_option = self._pytestconfig.getoption("--screenshot")
+        capture_screenshot = screenshot_option == "on" or (
+            failed and screenshot_option == "only-on-failure"
+        )
+        if capture_screenshot:
+            for index, screenshot in enumerate(self._screenshots):
+                human_readable_status = "failed" if failed else "finished"
+                screenshot_path = _build_artifact_test_folder(
+                    self._pytestconfig,
+                    self._request,
+                    f"test-{human_readable_status}-{index+1}.png",
+                )
+                os.makedirs(os.path.dirname(screenshot_path), exist_ok=True)
+                shutil.move(screenshot, screenshot_path)
+        else:
+            for screenshot in self._screenshots:
+                os.remove(screenshot)
+
+        if self._tracing_option == "on" or (
+            failed and self._tracing_option == "retain-on-failure"
+        ):
+            for index, trace in enumerate(self._traces):
+                trace_file_name = (
+                    "trace.zip" if len(self._traces) == 1 else f"trace-{index+1}.zip"
+                )
+                trace_path = _build_artifact_test_folder(
+                    self._pytestconfig, self._request, trace_file_name
+                )
+                os.makedirs(os.path.dirname(trace_path), exist_ok=True)
+                shutil.move(trace, trace_path)
+        else:
+            for trace in self._traces:
+                os.remove(trace)
+
+        video_option = self._pytestconfig.getoption("--video")
+        preserve_video = video_option == "on" or (
+            failed and video_option == "retain-on-failure"
+        )
+        if preserve_video:
+            for index, page in enumerate(self._all_pages):
+                video = page.video
+                if not video:
+                    continue
+                try:
+                    video_file_name = (
+                        "video.webm"
+                        if len(self._all_pages) == 1
+                        else f"video-{index+1}.webm"
+                    )
+                    video.save_as(
+                        path=_build_artifact_test_folder(
+                            self._pytestconfig, self._request, video_file_name
+                        )
+                    )
+                except Error:
+                    # Silent catch empty videos.
+                    pass
+
+    def on_did_create_browser_context(self, context: BrowserContext) -> None:
+        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),
+                screenshots=True,
+                snapshots=True,
+                sources=True,
+            )
+
+    def on_will_close_browser_context(self, context: BrowserContext) -> None:
+        if self._capture_trace:
+            trace_path = Path(artifacts_folder.name) / create_guid()
+            context.tracing.stop(path=trace_path)
+            self._traces.append(str(trace_path))
+        else:
+            context.tracing.stop()
+
+        if self._pytestconfig.getoption("--screenshot") in ["on", "only-on-failure"]:
+            for page in context.pages:
+                try:
+                    screenshot_path = Path(artifacts_folder.name) / create_guid()
+                    page.screenshot(
+                        timeout=5000,
+                        path=screenshot_path,
+                        full_page=self._pytestconfig.getoption(
+                            "--full-page-screenshot"
+                        ),
+                    )
+                    self._screenshots.append(str(screenshot_path))
+                except Error:
+                    pass
+
+
+def create_guid() -> str:
+    return hashlib.sha256(os.urandom(16)).hexdigest()