Skip to content

Commit

Permalink
feat: add multiple context support for tracing and video
Browse files Browse the repository at this point in the history
  • Loading branch information
mxschmitt committed Apr 15, 2024
1 parent be0be7c commit c0db2a6
Show file tree
Hide file tree
Showing 2 changed files with 327 additions and 72 deletions.
276 changes: 204 additions & 72 deletions pytest_playwright/pytest_playwright.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand All @@ -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-")


Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -228,93 +260,84 @@ 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] = []

) -> 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)

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:
_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
_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
return cast(CreateContextCallback, _new_context)

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")
Expand Down Expand Up @@ -419,3 +442,112 @@ 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._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"]

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, page in enumerate(self._all_pages):
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",
)
try:
page.screenshot(
timeout=5000,
path=screenshot_path,
full_page=self._pytestconfig.getoption(
"--full-page-screenshot"
),
)
except Error:
pass

# Close contexts which were not closed during the test (this will trigger Trace and Video generation)
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"
):
for index, trace in enumerate(self._traces):
retain_trace = self._capture_trace or failed
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
)
if retain_trace:
os.makedirs(os.path.dirname(trace_path), exist_ok=True)
shutil.move(trace, trace_path)
else:
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:
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),
screenshots=True,
snapshots=True,
sources=True,
)

def on_will_close_browser_context(self, context: BrowserContext) -> None:
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)
self._traces.append(str(trace_path))
else:
context.tracing.stop()


def create_guid() -> str:
return hashlib.sha256(os.urandom(16)).hexdigest()
Loading

0 comments on commit c0db2a6

Please sign in to comment.