Move profile handling code to a child class of SharedPageState

tools/perf/page_sets/pregenerated_profile_shared_state.py

+# Copyright 2016 The Chromium Authors. All rights reserved.
+# 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 tempfile
+import zipfile
+
+from catapult_base import cloud_storage
+
+from telemetry.page import shared_page_state
+
+
+class PregeneratedProfileSharedState(shared_page_state.SharedPageState):
+  def __init__(self, test, finder_options, story_set):
+    super(PregeneratedProfileSharedState, self).__init__(
+        test, finder_options, story_set)
+    self._unzipped_profile = None
+    self._migrated_profile = None
+    self._pregenerated_profile_archive_dir = None
+
+  def WillRunStory(self, page):
+    if self._ShouldDownloadPregeneratedProfileArchive():
+      self._DownloadPregeneratedProfileArchive()
+
+      if self._ShouldMigrateProfile():
+        self._MigratePregeneratedProfile()
+    super(PregeneratedProfileSharedState, self).WillRunStory(page)
+
+  def TearDownState(self):
+    if self._unzipped_profile:
+      shutil.rmtree(self._unzipped_profile)
+      self._unzipped_profile = None
+    if self._migrated_profile:
+      shutil.rmtree(self._migrated_profile)
+      self._migrated_profile = None
+    super(PregeneratedProfileSharedState, self).TearDownState()
+
+  def _ShouldDownloadPregeneratedProfileArchive(self):
+    """Whether to download a pre-generated profile archive."""
+    if not self._pregenerated_profile_archive_dir:
+      return False
+
+    if self._finder_options.browser_options.profile_dir:
+      logging.warning("Profile directory specified on command line: %s, this"
+                      "overrides the benchmark's default profile directory.",
+                      self._finder_options.browser_options.profile_dir)
+      return False
+
+    if self._possible_browser.IsRemote():
+      return False
+
+    return True
+
+  def _DownloadPregeneratedProfileArchive(self):
+    """Download and extract the profile directory archive if one exists.
+
+    On success, updates self._finder_options.browser_options.profile_dir with
+    the directory of the extracted profile.
+    """
+    try:
+      cloud_storage.GetIfChanged(self._pregenerated_profile_archive_dir,
+                                 cloud_storage.PUBLIC_BUCKET)
+    except (cloud_storage.CredentialsError,
+            cloud_storage.PermissionError) as e:
+      if os.path.exists(self._pregenerated_profile_archive_dir):
+        # If the profile directory archive exists, assume the user has their
+        # own local copy simply warn.
+        logging.warning('Could not download Profile archive: %s',
+                        self._pregenerated_profile_archive_dir)
+      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.',
+                      self._pregenerated_profile_archive_dir)
+        logging.error(str(e))
+        sys.exit(-1)
+
+    # Check to make sure the zip file exists.
+    if not os.path.isfile(self._pregenerated_profile_archive_dir):
+      raise Exception("Profile directory archive not downloaded: ",
+                      self._pregenerated_profile_archive_dir)
+
+    # The location to extract the profile into.
+    self._unzipped_profile = tempfile.mkdtemp()
+    profile_archive_path_basename = os.path.basename(
+        self._pregenerated_profile_archive_dir)
+    extracted_profile_dir_path = os.path.join(
+        self._unzipped_profile,
+        os.path.splitext(profile_archive_path_basename)[0])
+
+    # Unzip profile directory.
+    with zipfile.ZipFile(self._pregenerated_profile_archive_dir) as f:
+      try:
+        f.extractall(self._unzipped_profile)
+      except Exception as e:
+        # Cleanup any leftovers from unzipping.
+        shutil.rmtree(self._unzipped_profile)
+        logging.error("Error extracting profile directory zip file: %s", e)
+        sys.exit(-1)
+
+    if not os.path.exists(extracted_profile_dir_path):
+      raise Exception("Failed to extract profile: ",
+                      extracted_profile_dir_path)
+
+    # Run with freshly extracted profile directory.
+    logging.info("Using profile archive directory: %s",
+                 extracted_profile_dir_path)
+    self._finder_options.browser_options.profile_dir = (
+        extracted_profile_dir_path)
+
+  def _ShouldMigrateProfile(self):
+    return not self._migrated_profile
+
+  def _MigrateProfile(self, finder_options, found_browser,
+                      initial_profile, final_profile):
+    """Migrates a profile to be compatible with a newer version of Chrome.
+
+    Launching Chrome with the old profile will perform the migration.
+    """
+    # Save the current input and output profiles.
+    saved_input_profile = finder_options.browser_options.profile_dir
+    saved_output_profile = finder_options.browser_options.output_profile_path
+
+    # Set the input and output profiles.
+    finder_options.browser_options.profile_dir = initial_profile
+    finder_options.browser_options.output_profile_path = final_profile
+
+    # Launch the browser, then close it.
+    browser = found_browser.Create(finder_options)
+    browser.Close()
+
+    # Load the saved input and output profiles.
+    finder_options.browser_options.profile_dir = saved_input_profile
+    finder_options.browser_options.output_profile_path = saved_output_profile
+
+  def _MigratePregeneratedProfile(self):
+    """Migrates the pre-generated profile by launching Chrome with it.
+
+    On success, updates self._migrated_profile and
+    self._finder_options.browser_options.profile_dir with the directory of the
+    migrated profile.
+    """
+    self._migrated_profile = tempfile.mkdtemp()
+    logging.info("Starting migration of pre-generated profile to %s",
+                 self._migrated_profile)
+    pregenerated_profile = self._finder_options.browser_options.profile_dir
+
+    possible_browser = self._FindBrowser(self._finder_options)
+    self._MigrateProfile(self._finder_options, possible_browser,
+                         pregenerated_profile, self._migrated_profile)
+    self._finder_options.browser_options.profile_dir = self._migrated_profile
+    logging.info("Finished migration of pre-generated profile to %s",
+                 self._migrated_profile)