1#!/usr/bin/env python
2# Copyright 2017 The Chromium Authors. All rights reserved.
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5
6import argparse
7import json
8import logging
9import os
10import subprocess
11import sys
12
13
14def collect_task(
15    collect_cmd, merge_script, build_properties, merge_arguments,
16    task_output_dir, output_json):
17  """Collect and merge the results of a task.
18
19  This is a relatively thin wrapper script around a `swarming.py collect`
20  command and a subsequent results merge to ensure that the recipe system
21  treats them as a single step. The results merge can either be the default
22  one provided by results_merger or a python script provided as merge_script.
23
24  Args:
25    collect_cmd: The `swarming.py collect` command to run. Should not contain
26      a --task-output-dir argument.
27    merge_script: A merge/postprocessing script that should be run to
28      merge the results. This script will be invoked as
29
30        <merge_script> \
31            [--build-properties <string JSON>] \
32            [merge arguments...] \
33            --summary-json <summary json> \
34            -o <merged json path> \
35            <shard json>...
36
37      where the merge arguments are the contents of merge_arguments_json.
38    build_properties: A string containing build information to
39      pass to the merge script in JSON form.
40    merge_arguments: A string containing additional arguments to pass to
41      the merge script in JSON form.
42    task_output_dir: A path to a directory in which swarming will write the
43      output of the task, including a summary JSON and all of the individual
44      shard results.
45    output_json: A path to a JSON file to which the merged results should be
46      written. The merged results should be in the JSON Results File Format
47      (https://www.chromium.org/developers/the-json-test-results-format)
48      and may optionally contain a top level "links" field that may contain a
49      dict mapping link text to URLs, for a set of links that will be included
50      in the buildbot output.
51  Returns:
52    The exit code of collect_cmd or merge_cmd.
53  """
54  logging.debug('Using task_output_dir: %r', task_output_dir)
55  if os.path.exists(task_output_dir):
56    logging.warn('task_output_dir %r already exists!', task_output_dir)
57    existing_contents = []
58    try:
59      for p in os.listdir(task_output_dir):
60        existing_contents.append(os.path.join(task_output_dir, p))
61    except (OSError, IOError) as e:
62      logging.error('Error while examining existing task_output_dir: %s', e)
63
64    logging.warn('task_output_dir existing content: %r', existing_contents)
65
66  collect_cmd.extend(['--task-output-dir', task_output_dir])
67
68  logging.info('collect_cmd: %s', ' '.join(collect_cmd))
69  collect_result = subprocess.call(collect_cmd)
70  if collect_result != 0:
71    logging.warn('collect_cmd had non-zero return code: %s', collect_result)
72
73  task_output_dir_contents = []
74  try:
75    task_output_dir_contents.extend(
76        os.path.join(task_output_dir, p)
77        for p in os.listdir(task_output_dir))
78  except (OSError, IOError) as e:
79    logging.error('Error while processing task_output_dir: %s', e)
80
81  logging.debug('Contents of task_output_dir: %r', task_output_dir_contents)
82  if not task_output_dir_contents:
83    logging.warn(
84        'No files found in task_output_dir: %r',
85        task_output_dir)
86
87  task_output_subdirs = (
88      p for p in task_output_dir_contents
89      if os.path.isdir(p))
90  shard_json_files = [
91      os.path.join(subdir, 'output.json')
92      for subdir in task_output_subdirs]
93  extant_shard_json_files = [
94      f for f in shard_json_files if os.path.exists(f)]
95
96  if shard_json_files != extant_shard_json_files:
97    logging.warn(
98        'Expected output.json file missing: %r\nFound: %r\nExpected: %r\n',
99        set(shard_json_files) - set(extant_shard_json_files),
100        extant_shard_json_files,
101        shard_json_files)
102
103  if not extant_shard_json_files:
104    logging.warn(
105        'No shard json files found in task_output_dir: %r\nFound %r',
106        task_output_dir, task_output_dir_contents)
107
108  logging.debug('Found shard_json_files: %r', shard_json_files)
109
110  summary_json_file = os.path.join(task_output_dir, 'summary.json')
111
112  merge_result = 0
113
114  merge_cmd = [sys.executable, merge_script]
115  if build_properties:
116    merge_cmd.extend(('--build-properties', build_properties))
117  if os.path.exists(summary_json_file):
118    merge_cmd.extend(('--summary-json', summary_json_file))
119  else:
120    logging.warn('Summary json file missing: %r', summary_json_file)
121  if merge_arguments:
122    merge_cmd.extend(json.loads(merge_arguments))
123  merge_cmd.extend(('-o', output_json))
124  merge_cmd.extend(extant_shard_json_files)
125
126  logging.info('merge_cmd: %s', ' '.join(merge_cmd))
127  merge_result = subprocess.call(merge_cmd)
128  if merge_result != 0:
129    logging.warn('merge_cmd had non-zero return code: %s', merge_result)
130
131  if not os.path.exists(output_json):
132    logging.warn(
133        'merge_cmd did not create output_json file: %r', output_json)
134
135  return collect_result or merge_result
136
137
138def main():
139  parser = argparse.ArgumentParser()
140  parser.add_argument('--build-properties')
141  parser.add_argument('--merge-additional-args')
142  parser.add_argument('--merge-script', required=True)
143  parser.add_argument('--task-output-dir', required=True)
144  parser.add_argument('-o', '--output-json', required=True)
145  parser.add_argument('--verbose', action='store_true')
146  parser.add_argument('collect_cmd', nargs='+')
147
148  args = parser.parse_args()
149  if args.verbose:
150    logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
151
152  return collect_task(
153      args.collect_cmd,
154      args.merge_script, args.build_properties, args.merge_additional_args,
155      args.task_output_dir, args.output_json)
156
157
158if __name__ == '__main__':
159  sys.exit(main())
160