tp: Refactor diff_test script to standalone library
Bug:255535171
Change-Id: Ib4f12b7034feb9c8aa4fd183816ffc78a4cfd402
diff --git a/python/BUILD.gn b/python/BUILD.gn
index 6dcf9ed..06ebd9c 100644
--- a/python/BUILD.gn
+++ b/python/BUILD.gn
@@ -31,6 +31,14 @@
]
}
+perfetto_py_library("trace_processor_diff_tests") {
+ sources = [
+ "generators/diff_tests/runner.py",
+ "generators/diff_tests/testing.py",
+ "generators/diff_tests/utils.py",
+ ]
+}
+
perfetto_py_binary("trace_processor_py_example") {
sources = [ "example.py" ]
deps = [
diff --git a/python/generators/diff_tests/runner.py b/python/generators/diff_tests/runner.py
new file mode 100644
index 0000000..56465b3
--- /dev/null
+++ b/python/generators/diff_tests/runner.py
@@ -0,0 +1,357 @@
+#!/usr/bin/env python3
+# Copyright (C) 2023 The Android Open Source Project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import difflib
+import os
+import subprocess
+import sys
+import tempfile
+from dataclasses import dataclass
+from typing import List, Tuple
+import concurrent.futures
+
+from google.protobuf import text_format
+from python.generators.diff_tests.testing import DiffTest
+from python.generators.diff_tests.utils import (create_message_factory,
+ end_color, get_env, green, red,
+ yellow)
+from tools.proto_utils import serialize_python_trace, serialize_textproto_trace
+
+ROOT_DIR = os.path.dirname(
+ os.path.dirname(
+ os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
+
+
+# Performance result of running the test.
+@dataclass
+class PerfResult:
+ test_type: DiffTest.TestType
+ trace_path: str
+ query_path_or_metric: str
+ ingest_time_ns: int
+ real_time_ns: int
+
+ def __init__(self, test: DiffTest, perf_lines: List[str]):
+ self.test_type = test.type
+ self.trace_path = test.trace_path
+ self.query_path_or_metric = test.query_path
+
+ assert len(perf_lines) == 1
+ perf_numbers = perf_lines[0].split(',')
+
+ assert len(perf_numbers) == 2
+ self.ingest_time_ns = int(perf_numbers[0])
+ self.real_time_ns = int(perf_numbers[1])
+
+
+# Data gathered from running the test.
+@dataclass
+class TestResult:
+ test_type: DiffTest.TestType
+ input_name: str
+ trace: str
+ cmd: List[str]
+ expected: str
+ actual: str
+ passed: bool
+ stderr: str
+ exit_code: int
+ perf_lines: List[str]
+
+ def __init__(self, type: DiffTest.TestType, query: str, gen_trace_path: str,
+ cmd: List[str], expected_text: str, actual_text: str,
+ stderr: str, exit_code: int, perf_lines: List[str]) -> None:
+ self.test_type = type
+ self.input_name = query
+ self.trace = gen_trace_path
+ self.cmd = cmd
+ self.stderr = stderr
+ self.exit_code = exit_code
+ self.perf_lines = perf_lines
+ self.expected = expected_text
+ self.actual = actual_text
+
+ expected_content = expected_text.replace('\r\n', '\n')
+ actual_content = actual_text.replace('\r\n', '\n')
+ self.passed = (expected_content == actual_content)
+
+ def write_diff(self):
+ expected_lines = self.expected.splitlines(True)
+ actual_lines = self.actual.splitlines(True)
+ diff = difflib.unified_diff(
+ expected_lines, actual_lines, fromfile='expected', tofile='actual')
+ return "".join(list(diff))
+
+
+# Run a metrics based DiffTest.
+def run_metrics_test(test: DiffTest, trace_processor_path: str,
+ gen_trace_path: str,
+ metrics_message_factory) -> TestResult:
+ if test.expected_path:
+ with open(test.expected_path, 'r') as expected_file:
+ expected = expected_file.read()
+ else:
+ expected = test.blueprint.out
+ tmp_perf_file = tempfile.NamedTemporaryFile(delete=False)
+ json_output = os.path.basename(test.expected_path).endswith('.json.out')
+ cmd = [
+ trace_processor_path,
+ '--analyze-trace-proto-content',
+ '--crop-track-events',
+ '--run-metrics',
+ test.query_path,
+ '--metrics-output=%s' % ('json' if json_output else 'binary'),
+ '--perf-file',
+ tmp_perf_file.name,
+ gen_trace_path,
+ ]
+ tp = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.PIPE,
+ env=get_env(ROOT_DIR))
+ (stdout, stderr) = tp.communicate()
+
+ if json_output:
+ expected_text = expected
+ actual_text = stdout.decode('utf8')
+ else:
+ # Expected will be in text proto format and we'll need to parse it to
+ # a real proto.
+ expected_message = metrics_message_factory()
+ text_format.Merge(expected, expected_message)
+
+ # Actual will be the raw bytes of the proto and we'll need to parse it
+ # into a message.
+ actual_message = metrics_message_factory()
+ actual_message.ParseFromString(stdout)
+
+ # Convert both back to text format.
+ expected_text = text_format.MessageToString(expected_message)
+ actual_text = text_format.MessageToString(actual_message)
+
+ perf_lines = [line.decode('utf8') for line in tmp_perf_file.readlines()]
+ tmp_perf_file.close()
+ os.remove(tmp_perf_file.name)
+ return TestResult(test.type, test.query_path,
+ gen_trace_path, cmd, expected_text, actual_text,
+ stderr.decode('utf8'), tp.returncode, perf_lines)
+
+
+# Run a query based Diff Test.
+def run_query_test(test: DiffTest, trace_processor_path: str,
+ gen_trace_path: str) -> TestResult:
+ with open(test.expected_path, 'r') as expected_file:
+ expected = expected_file.read()
+ tmp_perf_file = tempfile.NamedTemporaryFile(delete=False)
+ cmd = [
+ trace_processor_path,
+ '--analyze-trace-proto-content',
+ '--crop-track-events',
+ '-q',
+ test.query_path if test.query_path else test.blueprint.query,
+ '--perf-file',
+ tmp_perf_file.name,
+ gen_trace_path,
+ ]
+ tp = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.PIPE,
+ env=get_env(ROOT_DIR))
+ (stdout, stderr) = tp.communicate()
+
+ perf_lines = [line.decode('utf8') for line in tmp_perf_file.readlines()]
+ tmp_perf_file.close()
+ os.remove(tmp_perf_file.name)
+
+ return TestResult(test.type, test.query_path, gen_trace_path, cmd, expected,
+ stdout.decode('utf8'), stderr.decode('utf8'), tp.returncode,
+ perf_lines)
+
+
+# Run a DiffTest
+def run_test(trace_descriptor_path: str, extension_descriptor_paths: List[str],
+ args: argparse.Namespace,
+ test: DiffTest) -> Tuple[str, bool, str, PerfResult]:
+ out_path = os.path.dirname(args.trace_processor)
+ if args.metrics_descriptor:
+ metrics_descriptor_paths = [args.metrics_descriptor]
+ else:
+ metrics_protos_path = os.path.join(out_path, 'gen', 'protos', 'perfetto',
+ 'metrics')
+ metrics_descriptor_paths = [
+ os.path.join(metrics_protos_path, 'metrics.descriptor'),
+ os.path.join(metrics_protos_path, 'chrome',
+ 'all_chrome_metrics.descriptor')
+ ]
+ metrics_message_factory = create_message_factory(
+ metrics_descriptor_paths, 'perfetto.protos.TraceMetrics')
+ result_str = ""
+ red_str = red(args.no_colors)
+ green_str = green(args.no_colors)
+ end_color_str = end_color(args.no_colors)
+ expected_path = test.expected_path
+ test_name = f"{test.name}"
+
+ if not os.path.exists(test.trace_path):
+ result_str += f"Trace file not found {test.trace_path}\n"
+ return test_name, False, result_str, None
+ elif not os.path.exists(expected_path):
+ result_str = f"Expected file not found {expected_path}"
+ return test_name, False, result_str, None
+
+ is_generated_trace = test.trace_path.endswith(
+ '.py') or test.trace_path.endswith('.textproto')
+ if test.trace_path.endswith('.py'):
+ gen_trace_file = tempfile.NamedTemporaryFile(delete=False)
+ serialize_python_trace(trace_descriptor_path, test.trace_path,
+ gen_trace_file)
+ gen_trace_path = os.path.realpath(gen_trace_file.name)
+ elif test.trace_path.endswith('.textproto'):
+ gen_trace_file = tempfile.NamedTemporaryFile(delete=False)
+ serialize_textproto_trace(trace_descriptor_path, extension_descriptor_paths,
+ test.trace_path, gen_trace_file)
+ gen_trace_path = os.path.realpath(gen_trace_file.name)
+ else:
+ gen_trace_file = None
+ gen_trace_path = test.trace_path
+
+ result_str += f"{yellow(args.no_colors)}[ RUN ]{end_color_str} "
+ result_str += f"{test_name}\n"
+
+ # We can't use delete=True here. When using that on Windows, the
+ # resulting file is opened in exclusive mode (in turn that's a subtle
+ # side-effect of the underlying CreateFile(FILE_ATTRIBUTE_TEMPORARY))
+ # and TP fails to open the passed path.
+ if test.type == DiffTest.TestType.QUERY:
+
+ if not os.path.exists(test.query_path):
+ result_str += f"Query file not found {test.query_path}"
+ return test_name, False, result_str, None
+
+ result = run_query_test(test, args.trace_processor, gen_trace_path)
+ elif test.type == DiffTest.TestType.METRIC:
+ result = run_metrics_test(test, args.trace_processor, gen_trace_path,
+ metrics_message_factory)
+ else:
+ assert False
+
+ if gen_trace_file:
+ if args.keep_input:
+ result_str += f"Saving generated input trace: {gen_trace_path}\n"
+ else:
+ gen_trace_file.close()
+ os.remove(gen_trace_path)
+
+ def write_cmdlines():
+ res = ""
+ if is_generated_trace:
+ res += 'Command to generate trace:\n'
+ res += 'tools/serialize_test_trace.py '
+ res += '--descriptor {} {} > {}\n'.format(
+ os.path.relpath(trace_descriptor_path, ROOT_DIR),
+ os.path.relpath(test.trace_path, ROOT_DIR),
+ os.path.relpath(gen_trace_path, ROOT_DIR))
+ res += f"Command line:\n{' '.join(result.cmd)}\n"
+ return res
+
+ if result.exit_code != 0 or not result.passed:
+ result_str += result.stderr
+
+ if result.exit_code == 0:
+ result_str += (
+ f"Expected did not match actual for trace "
+ f"{test.trace_path} and {result.test_type} {result.input_name}\n"
+ f"Expected file: {expected_path}\n")
+ result_str += write_cmdlines()
+ result_str += result.write_diff()
+ else:
+ result_str += write_cmdlines()
+
+ result_str += f"{red_str}[ FAILED ]{end_color_str} {test_name} "
+ result_str += f"{os.path.basename(test.trace_path)}\n"
+
+ if args.rebase:
+ if result.exit_code == 0:
+ result_str += f"Rebasing {expected_path}\n"
+ with open(expected_path, 'w') as f:
+ f.write(result.actual)
+ else:
+ result_str += f"Rebase failed for {expected_path} as query failed\n"
+
+ return test_name, False, result_str, None
+ else:
+ perf_result = PerfResult(test, result.perf_lines)
+
+ result_str += (f"{green_str}[ OK ]{end_color_str} {test.name} "
+ f"(ingest: {perf_result.ingest_time_ns / 1000000:.2f} ms "
+ f"query: {perf_result.real_time_ns / 1000000:.2f} ms)\n")
+ return test_name, True, result_str, perf_result
+
+
+# Run all DiffTests.
+def run_all_tests(trace_descriptor_path: str,
+ extension_descriptor_paths: List['str'],
+ args: argparse.Namespace, tests: List[DiffTest]
+ ) -> Tuple[List[str], List[PerfResult], List[str]]:
+ perf_data = []
+ test_failure = []
+ rebased = []
+ with concurrent.futures.ProcessPoolExecutor() as e:
+ fut = [
+ e.submit(run_test, trace_descriptor_path, extension_descriptor_paths,
+ args, test) for test in tests
+ ]
+ for res in concurrent.futures.as_completed(fut):
+ test_name, test_passed, res_str, perf_result = res.result()
+ sys.stderr.write(res_str)
+ if test_passed:
+ perf_data.append(perf_result)
+ else:
+ if args.rebase:
+ rebased.append(test_name)
+ test_failure.append(test_name)
+
+ return test_failure, perf_data, rebased
+
+
+# Load all DiffTests matching the patterns.
+def read_all_tests(query_metric_pattern, trace_pattern):
+ include_index_dir = os.path.join(ROOT_DIR, 'test', 'trace_processor')
+ tests = []
+
+ INCLUDE_PATH = os.path.join(ROOT_DIR, 'test', 'trace_processor')
+ sys.path.append(INCLUDE_PATH)
+ from include_index import fetch_all_diff_tests
+ sys.path.pop()
+ diff_tests = fetch_all_diff_tests(include_index_dir)
+
+ for test in diff_tests:
+ # Temporary assertion until string passing is supported.
+ if not (test.blueprint.is_out_file() and test.blueprint.is_query_file() and
+ test.blueprint.is_trace_file()):
+ raise AssertionError("Test parameters should be passed as files.")
+ if test.query_path and not query_metric_pattern.match(
+ os.path.basename(test.name)):
+ continue
+
+ if test.trace_path and not trace_pattern.match(
+ os.path.basename(test.trace_path)):
+ continue
+
+ tests.append(test)
+ return tests
diff --git a/python/generators/diff_tests/testing.py b/python/generators/diff_tests/testing.py
index 6eaba77..49d4ba0 100644
--- a/python/generators/diff_tests/testing.py
+++ b/python/generators/diff_tests/testing.py
@@ -20,11 +20,6 @@
from enum import Enum
-class TestType(Enum):
- QUERY = 1
- METRIC = 2
-
-
@dataclass
class Path:
filename: str
@@ -55,6 +50,10 @@
# script.
class DiffTest:
+ class TestType(Enum):
+ QUERY = 1
+ METRIC = 2
+
def __init__(self, name: str, blueprint: DiffTestBlueprint,
index_dir: str) -> None:
self.name = name
@@ -62,14 +61,14 @@
if blueprint.is_query_file():
if blueprint.query.filename.endswith('.sql'):
- self.type = TestType.QUERY
+ self.type = DiffTest.TestType.QUERY
self.query_path = os.path.abspath(
os.path.join(index_dir, blueprint.query.filename))
else:
- self.type = TestType.METRIC
+ self.type = DiffTest.TestType.METRIC
self.query_path = blueprint.query.filename
else:
- self.type = TestType.METRIC
+ self.type = DiffTest.TestType.METRIC
self.query_path = None
if blueprint.is_trace_file():
diff --git a/python/generators/diff_tests/utils.py b/python/generators/diff_tests/utils.py
new file mode 100644
index 0000000..d381f43
--- /dev/null
+++ b/python/generators/diff_tests/utils.py
@@ -0,0 +1,76 @@
+#!/usr/bin/env python3
+# Copyright (C) 2023 The Android Open Source Project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import sys
+import os
+import signal
+from typing import List
+
+from tools.proto_utils import create_message_factory
+
+USE_COLOR_CODES = sys.stderr.isatty()
+
+
+def red(no_colors):
+ return "\u001b[31m" if USE_COLOR_CODES and not no_colors else ""
+
+
+def green(no_colors):
+ return "\u001b[32m" if USE_COLOR_CODES and not no_colors else ""
+
+
+def yellow(no_colors):
+ return "\u001b[33m" if USE_COLOR_CODES and not no_colors else ""
+
+
+def end_color(no_colors):
+ return "\u001b[0m" if USE_COLOR_CODES and not no_colors else ""
+
+
+def get_env(root_dir):
+ env = {
+ 'PERFETTO_BINARY_PATH': os.path.join(root_dir, 'test', 'data'),
+ }
+ if sys.platform.startswith('linux'):
+ env['PATH'] = os.path.join(root_dir, 'buildtools', 'linux64', 'clang',
+ 'bin')
+ elif sys.platform.startswith('darwin'):
+ # Sadly, on macOS we need to check out the Android deps to get
+ # llvm symbolizer.
+ env['PATH'] = os.path.join(root_dir, 'buildtools', 'ndk', 'toolchains',
+ 'llvm', 'prebuilt', 'darwin-x86_64', 'bin')
+ elif sys.platform.startswith('win32'):
+ env['PATH'] = os.path.join(root_dir, 'buildtools', 'win', 'clang', 'bin')
+ return env
+
+
+def ctrl_c_handler(_num, _frame):
+ # Send a sigkill to the whole process group. Our process group looks like:
+ # - Main python interpreter running the main()
+ # - N python interpreters coming from ProcessPoolExecutor workers.
+ # - 1 trace_processor_shell subprocess coming from the subprocess.Popen().
+ # We don't need any graceful termination as the diff tests are stateless and
+ # don't write any file. Just kill them all immediately.
+ os.killpg(os.getpid(), signal.SIGKILL)
+
+
+def create_metrics_message_factory(metrics_descriptor_paths):
+ return create_message_factory(metrics_descriptor_paths,
+ 'perfetto.protos.TraceMetrics')
+
+
+def find_trace_descriptor(parent):
+ trace_protos_path = os.path.join(parent, 'gen', 'protos', 'perfetto', 'trace')
+ return os.path.join(trace_protos_path, 'trace.descriptor')
