Restructure uploader script so imports are less likely to break.

We need to modify sys.path before we import anything from catapult.
Unfortunately we need to modify it according to --outdir, so it needs
to happen at runtime rather than import time.

I try to split the script into a main which just sets up command line
args and sys.path and then imports the main script. This makes it less
likely that future maintainers will import something too early.

Bug: chromium:1029452
Change-Id: I16bf6257269ab8ab90dd74bff7880de8b5fb8071
Reviewed-on: https://webrtc-review.googlesource.com/c/src/+/170341
Commit-Queue: Patrik Höglund <phoglund@webrtc.org>
Reviewed-by: Mirko Bonadei <mbonadei@webrtc.org>
Cr-Commit-Position: refs/heads/master@{#30788}

2 files changed, 164 insertions, 153 deletions

diff --git a/tools_webrtc/perf/catapult_uploader.py b/tools_webrtc/perf/catapult_uploader.py
new file mode 100644
index 0000000000..d0b02f8c9c
--- /dev/null
+++ b/tools_webrtc/perf/catapult_uploader.py
@@ -0,0 +1,128 @@
+#!/usr/bin/env python
+# Copyright (c) 2020 The WebRTC project authors. All Rights Reserved.
+#
+# Use of this source code is governed by a BSD-style license
+# that can be found in the LICENSE file in the root of the source
+# tree. An additional intellectual property rights grant can be found
+# in the file PATENTS.  All contributing project authors may
+# be found in the AUTHORS file in the root of the source tree.
+
+
+import httplib2
+import json
+import subprocess
+import zlib
+
+from tracing.value import histogram_set
+from tracing.value.diagnostics import generic_set
+from tracing.value.diagnostics import reserved_infos
+
+
+def _GenerateOauthToken():
+  args = ['luci-auth', 'token']
+  p = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+  if p.wait() == 0:
+    output = p.stdout.read()
+    return output.strip()
+  else:
+    raise RuntimeError(
+        'Error generating authentication token.\nStdout: %s\nStderr:%s' %
+        (p.stdout.read(), p.stderr.read()))
+
+
+def _SendHistogramSet(url, histograms, oauth_token):
+  """Make a HTTP POST with the given JSON to the Performance Dashboard.
+
+  Args:
+    url: URL of Performance Dashboard instance, e.g.
+        "https://chromeperf.appspot.com".
+    histograms: a histogram set object that contains the data to be sent.
+    oauth_token: An oauth token to use for authorization.
+  """
+  headers = {'Authorization': 'Bearer %s' % oauth_token}
+
+  # TODO(https://crbug.com/1029452): HACKHACK
+  # Remove once we set bin bounds correctly in the proto writer.
+  dicts = histograms.AsDicts()
+  for d in dicts:
+    if 'name' in d:
+      d['allBins'] = [[1]]
+
+  serialized = json.dumps(dicts, indent=4)
+
+  if url.startswith('http://localhost'):
+    # The catapult server turns off compression in developer mode.
+    data = serialized
+  else:
+    data = zlib.compress(serialized)
+
+  print 'Sending %d bytes to %s.' % (len(data), url + '/add_histograms')
+
+  http = httplib2.Http()
+  response, content = http.request(url + '/add_histograms', method='POST',
+                                   body=data, headers=headers)
+  return response, content
+
+
+def _LoadHistogramSetFromProto(options):
+  hs = histogram_set.HistogramSet()
+  with options.input_results_file as f:
+    hs.ImportProto(f.read())
+
+  return hs
+
+
+def _AddBuildInfo(histograms, options):
+  common_diagnostics = {
+      reserved_infos.MASTERS: options.perf_dashboard_machine_group,
+      reserved_infos.BOTS: options.bot,
+      reserved_infos.POINT_ID: options.commit_position,
+      reserved_infos.BENCHMARKS: options.test_suite,
+      reserved_infos.WEBRTC_REVISIONS: str(options.webrtc_git_hash),
+      reserved_infos.BUILD_URLS: options.build_page_url,
+  }
+
+  for k, v in common_diagnostics.items():
+    histograms.AddSharedDiagnosticToAllHistograms(
+        k.name, generic_set.GenericSet([v]))
+
+
+def _DumpOutput(histograms, output_file):
+  with output_file:
+    json.dump(histograms.AsDicts(), output_file, indent=4)
+
+
+# TODO(https://crbug.com/1029452): Remove this once
+# https://chromium-review.googlesource.com/c/catapult/+/2094312 lands.
+def _HackSummaryOptions(histograms):
+  for histogram in histograms:
+    histogram.CustomizeSummaryOptions({
+      'avg': False,
+      'std': False,
+      'count': False,
+      'sum': False,
+      'min': False,
+      'max': False,
+      'nans': False,
+    })
+
+
+def UploadToDashboard(options):
+  histograms = _LoadHistogramSetFromProto(options)
+  _AddBuildInfo(histograms, options)
+  _HackSummaryOptions(histograms)
+
+  if options.output_json_file:
+    _DumpOutput(histograms, options.output_json_file)
+
+  oauth_token = _GenerateOauthToken()
+  response, content = _SendHistogramSet(
+      options.dashboard_url, histograms, oauth_token)
+
+  if response.status == 200:
+    print 'Received 200 from dashboard.'
+    return 0
+  else:
+    print('Upload failed with %d: %s\n\n%s' % (response.status, response.reason,
+                                               content))
+    return 1
diff --git a/tools_webrtc/perf/webrtc_dashboard_upload.py b/tools_webrtc/perf/webrtc_dashboard_upload.py
index 0c61400373..ed1b35e038 100644
--- a/tools_webrtc/perf/webrtc_dashboard_upload.py
+++ b/tools_webrtc/perf/webrtc_dashboard_upload.py
@@ -19,144 +19,8 @@ there. It needs to be here source-side.
 """
 
 import argparse
-import httplib2
-import json
 import os
 import sys
-import subprocess
-import zlib
-
-# We just yank the python scripts we require into the PYTHONPATH. You could also
-# imagine a solution where we use for instance protobuf:py_proto_runtime to copy
-# catapult and protobuf code to out/, but this approach is allowed by
-# convention. Fortunately neither catapult nor protobuf require any build rules
-# to be executed. We can't do this for the histogram proto stub though because
-# it's generated; see _LoadHistogramSetFromProto.
-#
-# It would be better if there was an equivalent to py_binary in GN, but there's
-# not.
-SCRIPT_DIR = os.path.dirname(os.path.realpath(__file__))
-CHECKOUT_ROOT = os.path.abspath(os.path.join(SCRIPT_DIR, os.pardir, os.pardir))
-sys.path.insert(0, os.path.join(CHECKOUT_ROOT, 'third_party', 'catapult',
-                                'tracing'))
-sys.path.insert(0, os.path.join(CHECKOUT_ROOT, 'third_party', 'protobuf',
-                                'python'))
-
-from tracing.value import histogram_set
-from tracing.value.diagnostics import generic_set
-from tracing.value.diagnostics import reserved_infos
-
-from google.protobuf import json_format
-
-
-def _GenerateOauthToken():
-  args = ['luci-auth', 'token']
-  p = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
-  if p.wait() == 0:
-    output = p.stdout.read()
-    return output.strip()
-  else:
-    raise RuntimeError(
-        'Error generating authentication token.\nStdout: %s\nStderr:%s' %
-        (p.stdout.read(), p.stderr.read()))
-
-
-def _SendHistogramSet(url, histograms, oauth_token):
-  """Make a HTTP POST with the given JSON to the Performance Dashboard.
-
-  Args:
-    url: URL of Performance Dashboard instance, e.g.
-        "https://chromeperf.appspot.com".
-    histograms: a histogram set object that contains the data to be sent.
-    oauth_token: An oauth token to use for authorization.
-  """
-  headers = {'Authorization': 'Bearer %s' % oauth_token}
-
-  # TODO(https://crbug.com/1029452): HACKHACK
-  # Remove once we set bin bounds correctly in the proto writer.
-  dicts = histograms.AsDicts()
-  for d in dicts:
-    if 'name' in d:
-      d['allBins'] = [[1]]
-
-  serialized = json.dumps(dicts, indent=4)
-
-  if url.startswith('http://localhost'):
-    # The catapult server turns off compression in developer mode.
-    data = serialized
-  else:
-    data = zlib.compress(serialized)
-
-  print 'Sending %d bytes to %s.' % (len(data), url + '/add_histograms')
-
-  http = httplib2.Http()
-  response, content = http.request(url + '/add_histograms', method='POST',
-                                   body=data, headers=headers)
-  return response, content
-
-
-def _LoadHistogramSetFromProto(options):
-  # The webrtc_dashboard_upload gn rule will build the protobuf stub for python,
-  # so put it in the path for this script before we attempt to import it.
-  histogram_proto_path = os.path.join(options.outdir, 'pyproto', 'tracing',
-                                      'tracing', 'proto')
-  sys.path.insert(0, histogram_proto_path)
-
-  # TODO(https://crbug.com/1029452): Get rid of this import hack once we can
-  # just hand the contents of input_results_file straight to the histogram set.
-  try:
-    import histogram_pb2
-  except ImportError:
-    raise ImportError('Could not find histogram_pb2. You need to build the '
-                      'webrtc_dashboard_upload target before invoking this '
-                      'script. Expected to find '
-                      'histogram_pb2 in %s.' % histogram_proto_path)
-
-  with options.input_results_file as f:
-    histograms = histogram_pb2.HistogramSet()
-    histograms.ParseFromString(f.read())
-
-  # TODO(https://crbug.com/1029452): Don't convert to JSON as a middle step once
-  # there is a proto de-serializer ready in catapult.
-  json_data = json.loads(json_format.MessageToJson(histograms))
-  hs = histogram_set.HistogramSet()
-  hs.ImportDicts(json_data)
-  return hs
-
-
-def _AddBuildInfo(histograms, options):
-  common_diagnostics = {
-      reserved_infos.MASTERS: options.perf_dashboard_machine_group,
-      reserved_infos.BOTS: options.bot,
-      reserved_infos.POINT_ID: options.commit_position,
-      reserved_infos.BENCHMARKS: options.test_suite,
-      reserved_infos.WEBRTC_REVISIONS: str(options.webrtc_git_hash),
-      reserved_infos.BUILD_URLS: options.build_page_url,
-  }
-
-  for k, v in common_diagnostics.items():
-    histograms.AddSharedDiagnosticToAllHistograms(
-        k.name, generic_set.GenericSet([v]))
-
-
-# TODO(https://crbug.com/1029452): Remove this once
-# https://chromium-review.googlesource.com/c/catapult/+/2094312 lands.
-def _HackSummaryOptions(histograms):
-  for histogram in histograms:
-    histogram.CustomizeSummaryOptions({
-      'avg': False,
-      'std': False,
-      'count': False,
-      'sum': False,
-      'min': False,
-      'max': False,
-      'nans': False,
-    })
-
-
-def _DumpOutput(histograms, output_file):
-  with output_file:
-    json.dump(histograms.AsDicts(), output_file, indent=4)
 
 
 def _CreateParser():
@@ -190,29 +54,48 @@ def _CreateParser():
   return parser
 
 
+def _ConfigurePythonPath(options):
+  # We just yank the python scripts we require into the PYTHONPATH. You could
+  # also imagine a solution where we use for instance protobuf:py_proto_runtime
+  # to copy catapult and protobuf code to out/. This is the convention in
+  # Chromium and WebRTC python scripts. We do need to build histogram_pb2
+  # however, so that's why we add out/ to sys.path below.
+  #
+  # It would be better if there was an equivalent to py_binary in GN, but
+  # there's not.
+  script_dir = os.path.dirname(os.path.realpath(__file__))
+  checkout_root = os.path.abspath(
+      os.path.join(script_dir, os.pardir, os.pardir))
+
+  sys.path.insert(0, os.path.join(checkout_root, 'third_party', 'catapult',
+                                  'tracing'))
+  sys.path.insert(0, os.path.join(checkout_root, 'third_party', 'protobuf',
+                                  'python'))
+
+  # The webrtc_dashboard_upload gn rule will build the protobuf stub for python,
+  # so put it in the path for this script before we attempt to import it.
+  histogram_proto_path = os.path.join(
+      options.outdir, 'pyproto', 'tracing', 'tracing', 'proto')
+  sys.path.insert(0, histogram_proto_path)
+
+  # Fail early in case the proto hasn't been built.
+  from tracing.proto import histogram_proto
+  if not histogram_proto.HAS_PROTO:
+    raise ImportError('Could not find histogram_pb2. You need to build the '
+                      'webrtc_dashboard_upload target before invoking this '
+                      'script. Expected to find '
+                      'histogram_pb2.py in %s.' % histogram_proto_path)
+
+
 def main(args):
   parser = _CreateParser()
   options = parser.parse_args(args)
 
-  histograms = _LoadHistogramSetFromProto(options)
-  _AddBuildInfo(histograms, options)
-  _HackSummaryOptions(histograms)
-
-  if options.output_json_file:
-    _DumpOutput(histograms, options.output_json_file)
-
-  oauth_token = _GenerateOauthToken()
-  response, content = _SendHistogramSet(
-      options.dashboard_url, histograms, oauth_token)
+  _ConfigurePythonPath(options)
 
-  if response.status == 200:
-    print 'Received 200 from dashboard.'
-    return 0
-  else:
-    print('Upload failed with %d: %s\n\n%s' % (response.status, response.reason,
-                                               content))
-    return 1
+  import catapult_uploader
 
+  return catapult_uploader.UploadToDashboard(options)
 
 if __name__ == '__main__':
   sys.exit(main(sys.argv[1:]))