xref: /aosp_15_r20/external/webrtc/tools_webrtc/gtest-parallel-wrapper.py (revision d9f758449e529ab9291ac668be2861e7a55c2422)
1#!/usr/bin/env vpython3
2
3# Copyright (c) 2016 The WebRTC project authors. All Rights Reserved.
4#
5# Use of this source code is governed by a BSD-style license
6# that can be found in the LICENSE file in the root of the source
7# tree. An additional intellectual property rights grant can be found
8# in the file PATENTS.  All contributing project authors may
9# be found in the AUTHORS file in the root of the source tree.
10
11# pylint: disable=invalid-name
12"""
13This script acts as an interface between the Chromium infrastructure and
14gtest-parallel, renaming options and translating environment variables into
15flags. Developers should execute gtest-parallel directly.
16
17In particular, this translates the GTEST_SHARD_INDEX and GTEST_TOTAL_SHARDS
18environment variables to the --shard_index and --shard_count flags, renames
19the --isolated-script-test-output flag to --dump_json_test_results,
20and interprets e.g. --workers=2x as 2 workers per core.
21
22Flags before '--' will be attempted to be understood as arguments to
23gtest-parallel. If gtest-parallel doesn't recognize the flag or the flag is
24after '--', the flag will be passed on to the test executable.
25
26--isolated-script-test-perf-output is renamed to
27--isolated_script_test_perf_output. The Android test runner needs the flag to
28be in the former form, but our tests require the latter, so this is the only
29place we can do it.
30
31If the --store-test-artifacts flag is set, an --output_dir must be also
32specified.
33
34The test artifacts will then be stored in a 'test_artifacts' subdirectory of the
35output dir, and will be compressed into a zip file once the test finishes
36executing.
37
38This is useful when running the tests in swarming, since the output directory
39is not known beforehand.
40
41For example:
42
43  gtest-parallel-wrapper.py some_test \
44      --some_flag=some_value \
45      --another_flag \
46      --output_dir=SOME_OUTPUT_DIR \
47      --store-test-artifacts
48      --isolated-script-test-output=SOME_DIR \
49      --isolated-script-test-perf-output=SOME_OTHER_DIR \
50      -- \
51      --foo=bar \
52      --baz
53
54Will be converted into:
55
56  vpython3 gtest-parallel \
57      --shard_index 0 \
58      --shard_count 1 \
59      --output_dir=SOME_OUTPUT_DIR \
60      --dump_json_test_results=SOME_DIR \
61      some_test \
62      -- \
63      --test_artifacts_dir=SOME_OUTPUT_DIR/test_artifacts \
64      --some_flag=some_value \
65      --another_flag \
66      --isolated-script-test-perf-output=SOME_OTHER_DIR \
67      --foo=bar \
68      --baz
69
70"""
71
72import argparse
73import collections
74import multiprocessing
75import os
76import shutil
77import subprocess
78import sys
79
80Args = collections.namedtuple(
81    'Args',
82    ['gtest_parallel_args', 'test_env', 'output_dir', 'test_artifacts_dir'])
83
84
85def _CatFiles(file_list, output_file_destination):
86  with open(output_file_destination, 'w') as output_file:
87    for filename in file_list:
88      with open(filename) as input_file:
89        output_file.write(input_file.read())
90      os.remove(filename)
91
92
93def _ParseWorkersOption(workers):
94  """Interpret Nx syntax as N * cpu_count. Int value is left as is."""
95  base = float(workers.rstrip('x'))
96  if workers.endswith('x'):
97    result = int(base * multiprocessing.cpu_count())
98  else:
99    result = int(base)
100  return max(result, 1)  # Sanitize when using e.g. '0.5x'.
101
102
103class ReconstructibleArgumentGroup:
104  """An argument group that can be converted back into a command line.
105
106  This acts like ArgumentParser.add_argument_group, but names of arguments added
107  to it are also kept in a list, so that parsed options from
108  ArgumentParser.parse_args can be reconstructed back into a command line (list
109  of args) based on the list of wanted keys."""
110
111  def __init__(self, parser, *args, **kwargs):
112    self._group = parser.add_argument_group(*args, **kwargs)
113    self._keys = []
114
115  def AddArgument(self, *args, **kwargs):
116    arg = self._group.add_argument(*args, **kwargs)
117    self._keys.append(arg.dest)
118
119  def RemakeCommandLine(self, options):
120    result = []
121    for key in self._keys:
122      value = getattr(options, key)
123      if value is True:
124        result.append('--%s' % key)
125      elif value is not None:
126        result.append('--%s=%s' % (key, value))
127    return result
128
129
130def ParseArgs(argv=None):
131  parser = argparse.ArgumentParser(argv)
132
133  gtest_group = ReconstructibleArgumentGroup(parser,
134                                             'Arguments to gtest-parallel')
135  # These options will be passed unchanged to gtest-parallel.
136  gtest_group.AddArgument('-d', '--output_dir')
137  gtest_group.AddArgument('-r', '--repeat')
138  # --isolated-script-test-output is used to upload results to the flakiness
139  # dashboard. This translation is made because gtest-parallel expects the flag
140  # to be called --dump_json_test_results instead.
141  gtest_group.AddArgument('--isolated-script-test-output',
142                          dest='dump_json_test_results')
143  gtest_group.AddArgument('--retry_failed')
144  gtest_group.AddArgument('--gtest_color')
145  gtest_group.AddArgument('--gtest_filter')
146  gtest_group.AddArgument('--gtest_also_run_disabled_tests',
147                          action='store_true',
148                          default=None)
149  gtest_group.AddArgument('--timeout')
150
151  # Syntax 'Nx' will be interpreted as N * number of cpu cores.
152  gtest_group.AddArgument('-w', '--workers', type=_ParseWorkersOption)
153
154  # Needed when the test wants to store test artifacts, because it doesn't
155  # know what will be the swarming output dir.
156  parser.add_argument('--store-test-artifacts', action='store_true')
157
158  parser.add_argument('executable')
159  parser.add_argument('executable_args', nargs='*')
160
161  options, unrecognized_args = parser.parse_known_args(argv)
162
163  executable_args = options.executable_args + unrecognized_args
164
165  if options.store_test_artifacts:
166    assert options.output_dir, (
167        '--output_dir must be specified for storing test artifacts.')
168    test_artifacts_dir = os.path.join(options.output_dir, 'test_artifacts')
169
170    executable_args.insert(0, '--test_artifacts_dir=%s' % test_artifacts_dir)
171  else:
172    test_artifacts_dir = None
173
174  gtest_parallel_args = gtest_group.RemakeCommandLine(options)
175
176  # GTEST_SHARD_INDEX and GTEST_TOTAL_SHARDS must be removed from the
177  # environment. Otherwise it will be picked up by the binary, causing a bug
178  # where only tests in the first shard are executed.
179  test_env = os.environ.copy()
180  gtest_shard_index = test_env.pop('GTEST_SHARD_INDEX', '0')
181  gtest_total_shards = test_env.pop('GTEST_TOTAL_SHARDS', '1')
182
183  gtest_parallel_args.insert(0, '--shard_index=%s' % gtest_shard_index)
184  gtest_parallel_args.insert(1, '--shard_count=%s' % gtest_total_shards)
185
186  gtest_parallel_args.append(options.executable)
187  if executable_args:
188    gtest_parallel_args += ['--'] + executable_args
189
190  return Args(gtest_parallel_args, test_env, options.output_dir,
191              test_artifacts_dir)
192
193
194def main():
195  webrtc_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
196  gtest_parallel_path = os.path.join(webrtc_root, 'third_party',
197                                     'gtest-parallel', 'gtest-parallel')
198
199  gtest_parallel_args, test_env, output_dir, test_artifacts_dir = ParseArgs()
200
201  command = [
202      sys.executable,
203      gtest_parallel_path,
204  ] + gtest_parallel_args
205
206  if output_dir and not os.path.isdir(output_dir):
207    os.makedirs(output_dir)
208  if test_artifacts_dir and not os.path.isdir(test_artifacts_dir):
209    os.makedirs(test_artifacts_dir)
210
211  print('gtest-parallel-wrapper: Executing command %s' % ' '.join(command))
212  sys.stdout.flush()
213
214  exit_code = subprocess.call(command, env=test_env, cwd=os.getcwd())
215
216  if output_dir:
217    for test_status in 'passed', 'failed', 'interrupted':
218      logs_dir = os.path.join(output_dir, 'gtest-parallel-logs', test_status)
219      if not os.path.isdir(logs_dir):
220        continue
221      logs = [os.path.join(logs_dir, log) for log in os.listdir(logs_dir)]
222      log_file = os.path.join(output_dir, '%s-tests.log' % test_status)
223      _CatFiles(logs, log_file)
224      os.rmdir(logs_dir)
225
226  if test_artifacts_dir:
227    shutil.make_archive(test_artifacts_dir, 'zip', test_artifacts_dir)
228    shutil.rmtree(test_artifacts_dir)
229
230  return exit_code
231
232
233if __name__ == '__main__':
234  sys.exit(main())
235