• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2024 The Chromium Authors
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5from __future__ import annotations
6
7import datetime as dt
8from typing import TYPE_CHECKING, List, Optional
9
10from crossbench.probes.probe import (Probe, ProbeConfigParser, ProbeContext,
11                                     ProbeMissingDataError)
12from crossbench.probes.result_location import ResultLocation
13from crossbench.probes.results import EmptyProbeResult, ProbeResult
14
15if TYPE_CHECKING:
16  from crossbench.browsers.browser import Viewport
17  from crossbench.env import HostEnvironment
18  from crossbench.path import AnyPath
19  from crossbench.runner.groups.browsers import BrowsersRunGroup
20  from crossbench.runner.groups.repetitions import RepetitionsRunGroup
21  from crossbench.runner.run import Run
22
23
24class ScreenshotProbe(Probe):
25  """
26  General-purpose Probe that collects screenshots.
27  """
28  NAME = "screenshot"
29  RESULT_LOCATION = ResultLocation.BROWSER
30  IMAGE_FORMAT = "png"
31
32  @classmethod
33  def config_parser(cls) -> ProbeConfigParser:
34    parser = super().config_parser()
35    # TODO: support interval-based screenshots
36    return parser
37
38  def _pre_check_viewport_size(self, env: HostEnvironment) -> None:
39    for browser in env.browsers:
40      viewport: Viewport = browser.viewport
41      if viewport.is_headless:
42        env.handle_warning(
43            f"Cannot take screenshots for headless browser: {browser}")
44      if viewport.x < 10 or viewport.y < 50:
45        env.handle_warning(
46            f"Viewport for '{browser}' might include toolbar: {viewport}")
47
48  def get_context(self, run: Run) -> ScreenshotProbeContext:
49    return ScreenshotProbeContext(self, run)
50
51  def merge_repetitions(self, group: RepetitionsRunGroup) -> ProbeResult:
52    # TODO: implement
53    return EmptyProbeResult()
54
55  def merge_browsers(self, group: BrowsersRunGroup) -> ProbeResult:
56    # TODO: implement
57    return EmptyProbeResult()
58
59
60class ScreenshotProbeContext(ProbeContext[ScreenshotProbe]):
61
62  def __init__(self, probe: ScreenshotProbe, run: Run) -> None:
63    super().__init__(probe, run)
64    self._results: List[AnyPath] = []
65
66  def get_default_result_path(self) -> AnyPath:
67    screenshot_dir = super().get_default_result_path()
68    self.browser_platform.mkdir(screenshot_dir)
69    return screenshot_dir
70
71  def start(self) -> None:
72    self.screenshot("start")
73
74  def start_story_run(self) -> None:
75    self.screenshot("start_story")
76
77  def stop_story_run(self) -> None:
78    self.screenshot("stop_story")
79
80  def stop(self) -> None:
81    self.screenshot("stop")
82
83  def screenshot(self, label: Optional[str] = None) -> None:
84    # TODO: support screen coordinates
85    if not label:
86      label = str(dt.datetime.now().strftime("%Y-%m-%d_%H%M%S"))
87    path = self.result_path / f"{label}.{ScreenshotProbe.IMAGE_FORMAT}"
88    # TODO: use the browser's implementation first which might be more portable
89    self.browser_platform.screenshot(path)
90    self._results.append(path)
91
92  def teardown(self) -> ProbeResult:
93    if not self.browser_platform.is_dir(self.result_path):
94      raise ProbeMissingDataError(
95          f"No screen shot found at: {self.result_path}")
96    return self.browser_result(file=tuple(self._results))
97