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