[Telemetry] Support downloading profiles from cloud storage

tools/telemetry/telemetry/test.py

Index: tools/telemetry/telemetry/test.py
diff --git a/tools/telemetry/telemetry/test.py b/tools/telemetry/telemetry/test.py
index 39003e2d310583c5f524c2c674469b5aa47fc8c4..ca2034047bee060498df08d2d962df55780fa5bc 100644
--- a/tools/telemetry/telemetry/test.py
+++ b/tools/telemetry/telemetry/test.py
@@ -2,11 +2,17 @@
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
+import logging
 import os
+import shutil
+import sys
+import zipfile
 
+from telemetry.core import browser_finder
 from telemetry.core import repeat_options
 from telemetry.core import util
 from telemetry.page import page_runner
+from telemetry.page import cloud_storage
 from telemetry.page import page_set
 from telemetry.page import page_test
 from telemetry.page import test_expectations
@@ -50,6 +56,8 @@ class Test(object):
     test.AddCommandLineOptions(parser)
     options.MergeDefaultValues(parser.get_default_values())
 
+    self._DownloadGeneratedProfileArchive(options)
+
     results = page_runner.Run(test, ps, expectations, options)
     results.PrintSummary()
     return len(results.failures) + len(results.errors)
@@ -62,6 +70,68 @@ class Test(object):
         getattr(options, 'pageset_repeat_iters', 1),
       )
 
+  def _DownloadGeneratedProfileArchive(self, options):
+    """Download and extract profile directory archive if one exists."""
+    archive_name = getattr(self, 'generated_profile_archive', None)
+
+    # If attribute not specified, nothing to do.
+    if not archive_name:
+      return
+
+    # If profile dire specfied on command line, nothing to do.
+    if options.browser_options.profile_dir:
+      logging.warning("Profile directory specified on command line: %s, this"
+          "overrides the benchmark's default profile directory.",
+          options.browser_options.profile_dir)
+      return
+
+    # Download profile directory from cloud storage.
+    test_data_dir = os.path.join(util.GetChromiumSrcDir(), 'tools', 'perf',
+        'generated_profiles',
+        browser_finder.TargetBrowserOS(options.browser_options))
+    generated_profile_archive_path = os.path.normpath(
+        os.path.join(test_data_dir, archive_name))
+
+    try:
+      cloud_storage.GetIfChanged(cloud_storage.PUBLIC_BUCKET,
+          generated_profile_archive_path)
+    except (cloud_storage.CredentialsError,
+            cloud_storage.PermissionError) as e:
+      if os.path.exists(generated_profile_archive_path):
+        # If the profile directory archive exists, assume the user has their
+        # own local copy simply warn.
+        logging.warning('Could not download Profile archive: %s',
+            generated_profile_archive_path)
+      else:
+        # If the archive profile directory doesn't exist, this is fatal.
+        logging.error('Can not run without required profile archive: %s. '
+                      'If you believe you have credentials, follow the '
+                      'instructions below.',
+                      generated_profile_archive_path)
+        logging.error(e)
+        sys.exit(1)
+
+    # Unzip profile directory.
+    extracted_profile_dir_path = (
+        os.path.splitext(generated_profile_archive_path)[0])
+    if not os.path.isfile(generated_profile_archive_path):
+      raise Exception("Profile directory archive not downloaded: ",
+          generated_profile_archive_path)
+    with zipfile.ZipFile(generated_profile_archive_path) as f:
+      try:
+        f.extractall(os.path.dirname(generated_profile_archive_path))
+      except e:
+        # Cleanup any leftovers from unzipping.
+        if os.path.exists(extracted_profile_dir_path):
+          shutil.rmtree(extracted_profile_dir_path)
+        logging.error("Error extracting profile directory zip file: %s", e)
+        sys.exit(1)
+
+    # Run with freshly extracted profile directory.
+    logging.info("Using profile archive directory: %s",
+        extracted_profile_dir_path)
+    options.browser_options.profile_dir = extracted_profile_dir_path
+
   def CreatePageSet(self, options):  # pylint: disable=W0613
     """Get the page set this test will run on.