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 7from typing import TYPE_CHECKING, Optional 8 9from crossbench.parse import ObjectParser 10from crossbench.probes.json import JsonResultProbe, JsonResultProbeContext 11from crossbench.probes.metric import MetricsMerger 12from crossbench.probes.probe import ProbeConfigParser, ProbeKeyT 13from crossbench.probes.result_location import ResultLocation 14 15if TYPE_CHECKING: 16 from crossbench.probes.results import ProbeResult 17 from crossbench.runner.actions import Actions 18 from crossbench.runner.groups.browsers import BrowsersRunGroup 19 from crossbench.runner.groups.stories import StoriesRunGroup 20 from crossbench.runner.run import Run 21 from crossbench.types import Json 22 23 24def parse_javascript(value: str) -> str: 25 # TODO: maybe add more sanity checks 26 return ObjectParser.non_empty_str(value, name="javascript") 27 28 29class JSProbe(JsonResultProbe): 30 """ 31 Probe for extracting arbitrary metrics using custom javascript code. 32 """ 33 NAME = "js" 34 RESULT_LOCATION = ResultLocation.LOCAL 35 IS_GENERAL_PURPOSE = True 36 37 @classmethod 38 def config_parser(cls) -> ProbeConfigParser: 39 parser = super().config_parser() 40 parser.add_argument( 41 "setup", 42 type=parse_javascript, 43 help=( 44 "Optional JavaScript code that is run immediately before a story. " 45 "This can be used for setting up some JS tracking code or patch " 46 "existing code for custom metric tracking.")) 47 parser.add_argument( 48 "js", 49 type=parse_javascript, 50 required=True, 51 help=("Required JavaScript code that is run immediately after " 52 "a story has finished. The code must return a JS object with " 53 "(nested) metric values (numbers).")) 54 return parser 55 56 def __init__(self, js: str, setup: Optional[str] = None) -> None: 57 super().__init__() 58 self._setup_js = setup 59 self._metric_js = js 60 61 @property 62 def setup_js(self) -> Optional[str]: 63 return self._setup_js 64 65 @property 66 def metric_js(self) -> str: 67 return self._metric_js 68 69 @property 70 def key(self) -> ProbeKeyT: 71 return super().key + ( 72 ("setup_js", self._setup_js), 73 ("metric_js", self._metric_js), 74 ) 75 76 def to_json(self, actions: Actions) -> Json: 77 data = actions.js(self._metric_js) 78 return ObjectParser.non_empty_dict(data, "JS metric data") 79 80 def get_context(self, run: Run) -> JSProbeContext: 81 return JSProbeContext(self, run) 82 83 def merge_stories(self, group: StoriesRunGroup) -> ProbeResult: 84 merged = MetricsMerger.merge_json_list( 85 story_group.results[self].json 86 for story_group in group.repetitions_groups) 87 return self.write_group_result(group, merged) 88 89 def merge_browsers(self, group: BrowsersRunGroup) -> ProbeResult: 90 return self.merge_browsers_json_list(group).merge( 91 self.merge_browsers_csv_list(group)) 92 93 94class JSProbeContext(JsonResultProbeContext[JSProbe]): 95 96 def start(self) -> None: 97 if setup_js := self.probe.setup_js: 98 with self.run.actions(f"Probe({self.probe.name}) setup") as actions: 99 actions.js(setup_js) 100