From 9e816118835ae04b37c701b1b2a0dadd821a910a Mon Sep 17 00:00:00 2001 From: Max Schmitt Date: Tue, 16 Apr 2024 20:36:25 +0200 Subject: [PATCH] review feedback --- pytest_playwright/pytest_playwright.py | 65 ++++++++++++++------------ 1 file changed, 36 insertions(+), 29 deletions(-) diff --git a/pytest_playwright/pytest_playwright.py b/pytest_playwright/pytest_playwright.py index 2d18432..db2cf0b 100644 --- a/pytest_playwright/pytest_playwright.py +++ b/pytest_playwright/pytest_playwright.py @@ -309,25 +309,30 @@ def new_context( browser_context_args: Dict, _artifacts_recorder: "ArtifactsRecorder", request: pytest.FixtureRequest, -) -> CreateContextCallback: +) -> 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] = [] def _new_context(**kwargs: Any) -> BrowserContext: context = browser.new_context(**browser_context_args, **kwargs) original_close = context.close - def close_wrapper(*args: Any, **kwargs: Any) -> None: + def _close_wrapper(*args: Any, **kwargs: Any) -> None: + contexts.remove(context) _artifacts_recorder.on_will_close_browser_context(context) original_close(*args, **kwargs) - context.close = close_wrapper + context.close = _close_wrapper + contexts.append(context) _artifacts_recorder.on_did_create_browser_context(context) return context - return cast(CreateContextCallback, _new_context) + yield cast(CreateContextCallback, _new_context) + for context in contexts.copy(): + context.close() @pytest.fixture @@ -452,8 +457,8 @@ def __init__( self._pytestconfig = pytestconfig self._playwright = playwright - self._open_contexts: BrowserContext = [] 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"] @@ -464,44 +469,34 @@ def did_finish_test(self, failed: bool) -> None: failed and screenshot_option == "only-on-failure" ) if capture_screenshot: - for index, page in enumerate(self._all_pages): + 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", ) - 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() + 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): - 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) + 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 ( @@ -528,7 +523,6 @@ def did_finish_test(self, failed: bool) -> None: 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( @@ -539,8 +533,6 @@ def on_did_create_browser_context(self, context: BrowserContext) -> None: ) 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) @@ -548,6 +540,21 @@ def on_will_close_browser_context(self, context: BrowserContext) -> None: 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()