Fix ability to capture archives in webpages_playback.py

tools/skp/webpages_playback.py

 #!/usr/bin/env python
 # Copyright (c) 2012 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.
 
 """Archives or replays webpages and creates SKPs in a Google Storage location.
 
 To archive webpages and store SKP files (archives should be rarely updated):
 
 cd ../buildbot/slave/skia_slave_scripts
@@ skipping 184 matching lines @@
 
     # Create the required local storage directories.
     self._CreateLocalStorageDirs()
 
     # Start the timer.
     start_time = time.time()
 
     # Loop through all page_sets.
     for page_set in self._page_sets:
 
-      page_set_basename = os.path.basename(page_set).split('.')[0] + '.json'
+      page_set_basename = os.path.basename(page_set).split('.')[0]
+      page_set_json_name = page_set_basename + '.json'
       wpr_data_file = page_set.split(os.path.sep)[-1].split('.')[0] + '_000.wpr'
+      page_set_dir = os.path.dirname(page_set)
 
       if self._record:
         # Create an archive of the specified webpages if '--record=True' is
         # specified.
         record_wpr_cmd = (
+          'PYTHONPATH=%s:$PYTHONPATH' % page_set_dir,
           'DISPLAY=%s' % X11_DISPLAY,
           os.path.join(self._telemetry_binaries_dir, 'record_wpr'),
           '--extra-browser-args=--disable-setuid-sandbox',
           '--browser=exact',
           '--browser-executable=%s' % self._browser_executable,
-          page_set
+          '%s_page_set' % page_set_basename,
+          '--page-set-base-dir=%s' % page_set_dir
         )
         for _ in range(RETRY_RECORD_WPR_COUNT):
           output = shell_utils.run(' '.join(record_wpr_cmd), shell=True)
           if RECORD_FAILURE_MSG in output:
             print output
           else:
             # Break out of the retry loop since there were no errors.
             break
         else:
           # If we get here then record_wpr did not succeed and thus did not
           # break out of the loop.
           raise Exception('record_wpr failed for page_set: %s' % page_set)
 
       else:
         if not self._skip_all_gs_access:
           # Get the webpages archive so that it can be replayed.
-          self._DownloadWebpagesArchive(wpr_data_file, page_set_basename)
+          self._DownloadWebpagesArchive(wpr_data_file, page_set_json_name)
 
-      page_set_name = os.path.basename(page_set).split('.')[0]
-      page_set_dir = os.path.dirname(page_set)
       run_benchmark_cmd = (
           'PYTHONPATH=%s:$PYTHONPATH' % page_set_dir,
           'DISPLAY=%s' % X11_DISPLAY,
           'timeout', '300',
           os.path.join(self._telemetry_binaries_dir, 'run_benchmark'),
           '--extra-browser-args=--disable-setuid-sandbox',
           '--browser=exact',
           '--browser-executable=%s' % self._browser_executable,
           SKP_BENCHMARK,
-          '--page-set-name=%s' % page_set_name,
+          '--page-set-name=%s' % page_set_basename,
           '--page-set-base-dir=%s' % page_set_dir,
           '--skp-outdir=%s' % TMP_SKP_DIR,
           '--also-run-disabled-tests'
       )
 
       for _ in range(RETRY_RUN_MEASUREMENT_COUNT):
         try:
           print '\n\n=======Capturing SKP of %s=======\n\n' % page_set
           shell_utils.run(' '.join(run_benchmark_cmd), shell=True)
         except shell_utils.CommandFailedException:
           # skpicture_printer sometimes fails with AssertionError but the
           # captured SKP is still valid. This is a known issue.
           pass
 
         if self._record:
           # Move over the created archive into the local webpages archive
           # directory.
           shutil.move(
               os.path.join(LOCAL_REPLAY_WEBPAGES_ARCHIVE_DIR, wpr_data_file),
               self._local_record_webpages_archive_dir)
           shutil.move(
               os.path.join(LOCAL_REPLAY_WEBPAGES_ARCHIVE_DIR,
-                           page_set_basename),
+                           page_set_json_name),
               self._local_record_webpages_archive_dir)
 
         # Rename generated SKP files into more descriptive names.
         try:
           self._RenameSkpFiles(page_set)
           # Break out of the retry loop since there were no errors.
           break
         except Exception:
           # There was a failure continue with the loop.
           traceback.print_exc()
@@ skipping 87 matching lines @@
       shutil.rmtree(site)
 
   def _CreateLocalStorageDirs(self):
     """Creates required local storage directories for this script."""
     for d in (self._local_record_webpages_archive_dir,
               self._local_skp_dir):
       if os.path.exists(d):
         shutil.rmtree(d)
       os.makedirs(d)
 
-  def _DownloadWebpagesArchive(self, wpr_data_file, page_set_basename):
+  def _DownloadWebpagesArchive(self, wpr_data_file, page_set_json_name):
     """Downloads the webpages archive and its required page set from GS."""
     wpr_source = posixpath.join(ROOT_PLAYBACK_DIR_NAME, 'webpages_archive',
                                 wpr_data_file)
     page_set_source = posixpath.join(ROOT_PLAYBACK_DIR_NAME,
                                      'webpages_archive',
-                                     page_set_basename)
+                                     page_set_json_name)
     gs = gs_utils.GSUtils()
     gs_bucket = self._dest_gsbase.lstrip(gs_utils.GS_PREFIX)
     if (gs.does_storage_object_exist(gs_bucket, wpr_source) and
         gs.does_storage_object_exist(gs_bucket, page_set_source)):
       gs.download_file(gs_bucket, wpr_source,
                        os.path.join(LOCAL_REPLAY_WEBPAGES_ARCHIVE_DIR,
                                     wpr_data_file))
       gs.download_file(gs_bucket, page_set_source,
                        os.path.join(LOCAL_REPLAY_WEBPAGES_ARCHIVE_DIR,
-                                    page_set_basename))
+                                    page_set_json_name))
     else:
       raise Exception('%s and %s do not exist in Google Storage!' % (
           wpr_source, page_set_source))
 
 
 if '__main__' == __name__:
   option_parser = optparse.OptionParser()
   option_parser.add_option(
       '', '--page_sets',
       help='Specifies the page sets to use to archive. Supports globs.',
@@ skipping 45 matching lines @@
       default=None)
   option_parser.add_option(
       '', '--non-interactive', action='store_true',
       help='Runs the script without any prompts. If this flag is specified and '
            '--skia_tools is specified then the debugger is not run.',
       default=False)
   options, unused_args = option_parser.parse_args()
 
   playback = SkPicturePlayback(options)
   sys.exit(playback.Run())