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