1*8975f5c5SAndroid Build Coastguard Worker#!/usr/bin/env vpython3 2*8975f5c5SAndroid Build Coastguard Worker 3*8975f5c5SAndroid Build Coastguard Worker# Copyright 2024 The Chromium Authors 4*8975f5c5SAndroid Build Coastguard Worker# Use of this source code is governed by a BSD-style license that can be 5*8975f5c5SAndroid Build Coastguard Worker# found in the LICENSE file. 6*8975f5c5SAndroid Build Coastguard Worker""" The module to create and manage measures using in the process. """ 7*8975f5c5SAndroid Build Coastguard Worker 8*8975f5c5SAndroid Build Coastguard Workerimport json 9*8975f5c5SAndroid Build Coastguard Workerimport os 10*8975f5c5SAndroid Build Coastguard Worker 11*8975f5c5SAndroid Build Coastguard Workerfrom google.protobuf import any_pb2 12*8975f5c5SAndroid Build Coastguard Workerfrom google.protobuf.json_format import MessageToDict 13*8975f5c5SAndroid Build Coastguard Worker 14*8975f5c5SAndroid Build Coastguard Workerfrom average import Average 15*8975f5c5SAndroid Build Coastguard Workerfrom count import Count 16*8975f5c5SAndroid Build Coastguard Workerfrom data_points import DataPoints 17*8975f5c5SAndroid Build Coastguard Workerfrom measure import Measure 18*8975f5c5SAndroid Build Coastguard Workerfrom metric import Metric 19*8975f5c5SAndroid Build Coastguard Workerfrom time_consumption import TimeConsumption 20*8975f5c5SAndroid Build Coastguard Worker 21*8975f5c5SAndroid Build Coastguard Worker# This is used as the key when being uploaded to ResultDB via result_sink 22*8975f5c5SAndroid Build Coastguard Worker# and shouldn't be changed 23*8975f5c5SAndroid Build Coastguard WorkerTEST_SCRIPT_METRICS_KEY = 'test_script_metrics' 24*8975f5c5SAndroid Build Coastguard Worker 25*8975f5c5SAndroid Build Coastguard Worker# The file name is used as the key when being loaded into the ResultDB and 26*8975f5c5SAndroid Build Coastguard Worker# shouldn't be changed. 27*8975f5c5SAndroid Build Coastguard WorkerTEST_SCRIPT_METRICS_JSONPB_FILENAME = f'{TEST_SCRIPT_METRICS_KEY}.jsonpb' 28*8975f5c5SAndroid Build Coastguard Worker 29*8975f5c5SAndroid Build Coastguard Worker_metric = Metric() 30*8975f5c5SAndroid Build Coastguard Worker 31*8975f5c5SAndroid Build Coastguard Worker 32*8975f5c5SAndroid Build Coastguard Workerdef _create_name(*name_pieces: str) -> str: 33*8975f5c5SAndroid Build Coastguard Worker if len(name_pieces) == 0: 34*8975f5c5SAndroid Build Coastguard Worker raise ValueError('Need at least one name piece.') 35*8975f5c5SAndroid Build Coastguard Worker return '/'.join(list(name_pieces)) 36*8975f5c5SAndroid Build Coastguard Worker 37*8975f5c5SAndroid Build Coastguard Worker 38*8975f5c5SAndroid Build Coastguard Workerdef _register(m: Measure) -> Measure: 39*8975f5c5SAndroid Build Coastguard Worker _metric.register(m) 40*8975f5c5SAndroid Build Coastguard Worker return m 41*8975f5c5SAndroid Build Coastguard Worker 42*8975f5c5SAndroid Build Coastguard Worker 43*8975f5c5SAndroid Build Coastguard Workerdef average(*name_pieces: str) -> Average: 44*8975f5c5SAndroid Build Coastguard Worker return _register(Average(_create_name(*name_pieces))) 45*8975f5c5SAndroid Build Coastguard Worker 46*8975f5c5SAndroid Build Coastguard Worker 47*8975f5c5SAndroid Build Coastguard Workerdef count(*name_pieces: str) -> Count: 48*8975f5c5SAndroid Build Coastguard Worker return _register(Count(_create_name(*name_pieces))) 49*8975f5c5SAndroid Build Coastguard Worker 50*8975f5c5SAndroid Build Coastguard Worker 51*8975f5c5SAndroid Build Coastguard Workerdef data_points(*name_pieces: str) -> DataPoints: 52*8975f5c5SAndroid Build Coastguard Worker return _register(DataPoints(_create_name(*name_pieces))) 53*8975f5c5SAndroid Build Coastguard Worker 54*8975f5c5SAndroid Build Coastguard Worker 55*8975f5c5SAndroid Build Coastguard Workerdef time_consumption(*name_pieces: str) -> TimeConsumption: 56*8975f5c5SAndroid Build Coastguard Worker return _register(TimeConsumption(_create_name(*name_pieces))) 57*8975f5c5SAndroid Build Coastguard Worker 58*8975f5c5SAndroid Build Coastguard Worker 59*8975f5c5SAndroid Build Coastguard Workerdef tag(*args: str) -> None: 60*8975f5c5SAndroid Build Coastguard Worker """Adds a tag to the Metric to tag the final results; see Metric for details. 61*8975f5c5SAndroid Build Coastguard Worker """ 62*8975f5c5SAndroid Build Coastguard Worker _metric.tag(*args) 63*8975f5c5SAndroid Build Coastguard Worker 64*8975f5c5SAndroid Build Coastguard Worker 65*8975f5c5SAndroid Build Coastguard Workerdef clear() -> None: 66*8975f5c5SAndroid Build Coastguard Worker """Clears all the registered Measures.""" 67*8975f5c5SAndroid Build Coastguard Worker _metric.clear() 68*8975f5c5SAndroid Build Coastguard Worker 69*8975f5c5SAndroid Build Coastguard Worker 70*8975f5c5SAndroid Build Coastguard Workerdef size() -> int: 71*8975f5c5SAndroid Build Coastguard Worker """Gets the current size of registered Measures.""" 72*8975f5c5SAndroid Build Coastguard Worker return _metric.size() 73*8975f5c5SAndroid Build Coastguard Worker 74*8975f5c5SAndroid Build Coastguard Workerdef to_dict() -> dict: 75*8975f5c5SAndroid Build Coastguard Worker """Converts all the registered Measures to a dict. 76*8975f5c5SAndroid Build Coastguard Worker 77*8975f5c5SAndroid Build Coastguard Worker The records are wrapped in protobuf Any message before exported as dict 78*8975f5c5SAndroid Build Coastguard Worker so that an additional key "@type" is included. 79*8975f5c5SAndroid Build Coastguard Worker """ 80*8975f5c5SAndroid Build Coastguard Worker any_msg = any_pb2.Any() 81*8975f5c5SAndroid Build Coastguard Worker any_msg.Pack(_metric.dump()) 82*8975f5c5SAndroid Build Coastguard Worker return MessageToDict(any_msg, preserving_proto_field_name=True) 83*8975f5c5SAndroid Build Coastguard Worker 84*8975f5c5SAndroid Build Coastguard Worker 85*8975f5c5SAndroid Build Coastguard Workerdef to_json() -> str: 86*8975f5c5SAndroid Build Coastguard Worker """Converts all the registered Measures to a json str.""" 87*8975f5c5SAndroid Build Coastguard Worker return json.dumps(to_dict(), sort_keys=True, indent=2) 88*8975f5c5SAndroid Build Coastguard Worker 89*8975f5c5SAndroid Build Coastguard Worker# TODO(crbug.com/343242386): May need to implement a lock and reset logic to 90*8975f5c5SAndroid Build Coastguard Worker# clear in-memory data and lock the instance to block further operations and 91*8975f5c5SAndroid Build Coastguard Worker# avoid accidentally accumulating data which won't be published at all. 92*8975f5c5SAndroid Build Coastguard Workerdef dump(dir_path: str) -> None: 93*8975f5c5SAndroid Build Coastguard Worker """Dumps the metric data into test_script_metrics.jsonpb in the |path|.""" 94*8975f5c5SAndroid Build Coastguard Worker os.makedirs(dir_path, exist_ok=True) 95*8975f5c5SAndroid Build Coastguard Worker with open(os.path.join(dir_path, TEST_SCRIPT_METRICS_JSONPB_FILENAME), 96*8975f5c5SAndroid Build Coastguard Worker 'w', 97*8975f5c5SAndroid Build Coastguard Worker encoding='utf-8') as wf: 98*8975f5c5SAndroid Build Coastguard Worker wf.write(to_json()) 99