• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1#!/usr/bin/env python3
2# Copyright (C) 2023 The Android Open Source Project
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#      http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15
16import inspect
17import os
18from dataclasses import dataclass
19from typing import Any, Dict, List, Optional, Union, Callable, Tuple
20from enum import Enum
21import re
22
23from google.protobuf import text_format
24
25TestName = str
26
27
28@dataclass
29class Path:
30  filename: str
31
32
33@dataclass
34class DataPath(Path):
35  filename: str
36
37
38@dataclass
39class Metric:
40  name: str
41
42
43@dataclass
44class MetricV2SpecTextproto:
45  contents: str
46
47
48@dataclass
49class Json:
50  contents: str
51
52
53@dataclass
54class Csv:
55  contents: str
56
57
58@dataclass
59class TextProto:
60  contents: str
61
62
63@dataclass
64class BinaryProto:
65  message_type: str
66  contents: str
67  # Comparing protos is tricky. For example, repeated fields might be written in
68  # any order. To help with that you can specify a `post_processing` function
69  # that will be called with the actual proto message object before converting
70  # it to text representation and doing the comparison with `contents`. This
71  # gives us a chance to e.g. sort messages in a repeated field.
72  post_processing: Callable = text_format.MessageToString
73
74
75@dataclass
76class Systrace:
77  contents: str
78
79
80class TraceInjector:
81  '''Injects fields into trace packets before test starts.
82
83  TraceInjector can be used within a DiffTestBlueprint to selectively inject
84  fields to trace packets containing specific data types. For example:
85
86    DiffTestBlueprint(
87        trace=...,
88        trace_modifier=TraceInjector('ftrace_events',
89                                     'sys_stats',
90                                     'process_tree',
91                                     {'machine_id': 1001},
92                                     trusted_uid=123)
93        query=...,
94        out=...)
95
96  packet_data_types: Data types to target for injection ('ftrace_events',
97  'sys_stats', 'process_tree')
98  injected_fields: Fields and their values to inject into matching packets
99  ({'machine_id': 1001}, trusted_uid=123).
100  '''
101
102  def __init__(self, packet_data_types: List[str], injected_fields: Dict[str,
103                                                                         Any]):
104    self.packet_data_types = packet_data_types
105    self.injected_fields = injected_fields
106
107  def inject(self, proto):
108    for p in proto.packet:
109      for f in self.packet_data_types:
110        if p.HasField(f):
111          for k, v, in self.injected_fields.items():
112            setattr(p, k, v)
113          continue
114
115
116class TestType(Enum):
117  QUERY = 1
118  METRIC = 2
119  METRIC_V2 = 3
120
121
122# Blueprint for running the diff test. 'query' is being run over data from the
123# 'trace 'and result will be compared to the 'out. Each test (function in class
124# inheriting from TestSuite) returns a DiffTestBlueprint.
125@dataclass
126class DiffTestBlueprint:
127
128  trace: Union[Path, DataPath, Json, Systrace, TextProto]
129  query: Union[str, Path, DataPath, Metric, MetricV2SpecTextproto]
130  out: Union[Path, DataPath, Json, Csv, TextProto, BinaryProto]
131  trace_modifier: Union[TraceInjector, None] = None
132  register_files_dir: Optional[DataPath] = None
133
134  def is_trace_file(self):
135    return isinstance(self.trace, Path)
136
137  def is_trace_textproto(self):
138    return isinstance(self.trace, TextProto)
139
140  def is_trace_json(self):
141    return isinstance(self.trace, Json)
142
143  def is_trace_systrace(self):
144    return isinstance(self.trace, Systrace)
145
146  def is_query_file(self):
147    return isinstance(self.query, Path)
148
149  def is_metric(self):
150    return isinstance(self.query, Metric)
151
152  def is_metric_v2(self):
153    return isinstance(self.query, MetricV2SpecTextproto)
154
155  def is_out_file(self):
156    return isinstance(self.out, Path)
157
158  def is_out_json(self):
159    return isinstance(self.out, Json)
160
161  def is_out_texproto(self):
162    return isinstance(self.out, TextProto)
163
164  def is_out_binaryproto(self):
165    return isinstance(self.out, BinaryProto)
166
167  def is_out_csv(self):
168    return isinstance(self.out, Csv)
169
170
171# Description of a diff test. Created in `fetch_diff_tests()` in
172# TestSuite: each test (function starting with `test_`) returns
173# DiffTestBlueprint and function name is a TestCase name. Used by diff test
174# script.
175class TestCase:
176
177  def __init__(
178      self,
179      name: str,
180      blueprint: DiffTestBlueprint,
181      index_dir: str,
182      test_data_dir: str,
183  ) -> None:
184    self.name = name
185    self.blueprint = blueprint
186    self.index_dir = index_dir
187    self.test_data_dir = test_data_dir
188
189    if blueprint.is_metric():
190      self.type = TestType.METRIC
191    elif blueprint.is_metric_v2():
192      self.type = TestType.METRIC_V2
193    else:
194      self.type = TestType.QUERY
195
196    self.query_path = self.__get_query_path()
197    self.trace_path = self.__get_trace_path()
198    self.expected_path = self.__get_expected_path()
199    self.expected_str = self.__get_expected_str()
200    self.register_files_dir = self.__get_register_files_dir()
201
202  def __get_query_path(self) -> Optional[str]:
203    if not self.blueprint.is_query_file():
204      return None
205
206    if isinstance(self.blueprint.query, DataPath):
207      path = os.path.join(self.test_data_dir, self.blueprint.query.filename)
208    else:
209      assert isinstance(self.blueprint.query, Path)
210      path = os.path.abspath(
211          os.path.join(self.index_dir, self.blueprint.query.filename))
212
213    if not os.path.exists(path):
214      raise AssertionError(
215          f"Query file ({path}) for test '{self.name}' does not exist.")
216    return path
217
218  def __get_trace_path(self) -> Optional[str]:
219    if not self.blueprint.is_trace_file():
220      return None
221
222    if isinstance(self.blueprint.trace, DataPath):
223      path = os.path.join(self.test_data_dir, self.blueprint.trace.filename)
224    else:
225      assert isinstance(self.blueprint.trace, Path)
226      path = os.path.abspath(
227          os.path.join(self.index_dir, self.blueprint.trace.filename))
228
229    if not os.path.exists(path):
230      raise AssertionError(
231          f"Trace file ({path}) for test '{self.name}' does not exist.")
232    return path
233
234  def __get_expected_path(self) -> Optional[str]:
235    if not self.blueprint.is_out_file():
236      return None
237
238    if isinstance(self.blueprint.out, DataPath):
239      path = os.path.join(self.test_data_dir, self.blueprint.out.filename)
240    else:
241      assert isinstance(self.blueprint.out, Path)
242      path = os.path.abspath(
243          os.path.join(self.index_dir, self.blueprint.out.filename))
244
245    if not os.path.exists(path):
246      raise AssertionError(
247          f"Out file ({path}) for test '{self.name}' does not exist.")
248    return path
249
250  def __get_register_files_dir(self) -> Optional[str]:
251    if not self.blueprint.register_files_dir:
252      return None
253
254    path = os.path.join(self.test_data_dir,
255                        self.blueprint.register_files_dir.filename)
256
257    if not os.path.exists(path):
258      raise AssertionError(
259          f"Out file ({path}) for test '{self.name}' does not exist.")
260    return path
261
262  def __get_expected_str(self) -> str:
263    if self.blueprint.is_out_file():
264      assert self.expected_path
265      with open(self.expected_path, 'r') as expected_file:
266        return expected_file.read()
267    assert isinstance(self.blueprint.out, (
268        TextProto,
269        Json,
270        Csv,
271        BinaryProto,
272        Systrace,
273    ))
274    return self.blueprint.out.contents
275
276  # Verifies that the test should be in test suite. If False, test will not be
277  # executed.
278  def validate(self, name_filter: str):
279    query_metric_pattern = re.compile(name_filter)
280    return bool(query_metric_pattern.match(os.path.basename(self.name)))
281
282
283# str.removeprefix is available in Python 3.9+, but the Perfetto CI runs on
284# older versions.
285def removeprefix(s: str, prefix: str):
286  if s.startswith(prefix):
287    return s[len(prefix):]
288  return s
289
290# Virtual class responsible for fetching diff tests.
291# All functions with name starting with `test_` have to return
292# DiffTestBlueprint and function name is a test name. All DiffTestModules have
293# to be included in `test/diff_tests/trace_processor/include_index.py`.
294# `fetch` function should not be overwritten.
295class TestSuite:
296
297  def __init__(
298      self,
299      include_index_dir: str,
300      test_data_dir: str = os.path.abspath(
301          os.path.join(__file__, '../../../../test/data'))
302  ) -> None:
303    # The last path in the module is the module name itself, which is not a part
304    # of the directory. The first part is "diff_tests.", but it is not present
305    # when running difftests from Chrome, so we strip it conditionally.
306    self.dir_name = '/'.join(
307        removeprefix(self.__class__.__module__, 'diff_tests.').split('.')[:-1])
308    self.index_dir = os.path.join(include_index_dir, self.dir_name)
309    self.class_name = self.__class__.__name__
310    self.test_data_dir = test_data_dir
311
312  def __test_name(self, method_name):
313    return f"{self.class_name}:{method_name.split('test_',1)[1]}"
314
315  def fetch(self) -> List['TestCase']:
316    attrs = (getattr(self, name) for name in dir(self))
317    methods = [attr for attr in attrs if inspect.ismethod(attr)]
318    return [
319        TestCase(
320            self.__test_name(method.__name__), method(), self.index_dir,
321            self.test_data_dir)
322        for method in methods
323        if method.__name__.startswith('test_')
324    ]
325
326
327def PrintProfileProto(profile):
328  locations = {l.id: l for l in profile.location}
329  functions = {f.id: f for f in profile.function}
330  samples = []
331  # Strips trailing annotations like (.__uniq.1657) from the function name.
332  filter_fname = lambda x: re.sub(' [(\[].*?uniq.*?[)\]]$', '', x)
333  for s in profile.sample:
334    stack = []
335    for location in [locations[id] for id in s.location_id]:
336      for function in [functions[l.function_id] for l in location.line]:
337        stack.append("{name} ({address})".format(
338            name=filter_fname(profile.string_table[function.name]),
339            address=hex(location.address)))
340      if len(location.line) == 0:
341        stack.append("({address})".format(address=hex(location.address)))
342    samples.append('Sample:\nValues: {values}\nStack:\n{stack}'.format(
343        values=', '.join(map(str, s.value)), stack='\n'.join(stack)))
344  return '\n\n'.join(sorted(samples)) + '\n'
345