Refactor of record_wpr.py

tools/telemetry/telemetry/page/record_wpr.py

Index: tools/telemetry/telemetry/page/record_wpr.py
diff --git a/tools/telemetry/telemetry/page/record_wpr.py b/tools/telemetry/telemetry/page/record_wpr.py
index 9491e4d9a8280ef57263c2bb12f2e42089b6dd72..561a7f9d85e95c960bf8b011499dbb4d6415fa92 100755
--- a/tools/telemetry/telemetry/page/record_wpr.py
+++ b/tools/telemetry/telemetry/page/record_wpr.py
@@ -18,20 +18,14 @@ from telemetry.page import page_set
 from telemetry.page import page_test
 from telemetry.page import profile_creator
 from telemetry.page import test_expectations
-from telemetry.page.actions import action_runner as action_runner_module
 from telemetry.results import page_measurement_results
 
 
-class RecordPage(page_test.PageTest):  # pylint: disable=W0223
-  def __init__(self, measurements):
-    # This class overwrites PageTest.Run, so that the test method name is not
-    # really used (except for throwing an exception if it doesn't exist).
-    super(RecordPage, self).__init__('Run')
-    self._action_names = set(
-        [measurement().action_name_to_run
-         for measurement in measurements.values()
-         if measurement().action_name_to_run])
-    self.test = None
+class RecorderPageTest(page_test.PageTest):  # pylint: disable=W0223
+  def __init__(self, action_names):
+    super(RecorderPageTest, self).__init__()
+    self._action_names = action_names
+    self.page_test = None
 
   def CanRunForPage(self, page):
     return page.url.startswith('http')
@@ -39,14 +33,25 @@ class RecordPage(page_test.PageTest):  # pylint: disable=W0223
   def WillNavigateToPage(self, page, tab):
     """Override to ensure all resources are fetched from network."""
     tab.ClearCache(force=False)
-    if self.test:
-      self.test.options = self.options
-      self.test.WillNavigateToPage(page, tab)
+    if self.page_test:
+      self.page_test.options = self.options
+      self.page_test.WillNavigateToPage(page, tab)
 
   def DidNavigateToPage(self, page, tab):
-    """Forward the call to the test."""
-    if self.test:
-      self.test.DidNavigateToPage(page, tab)
+    if self.page_test:
+      self.page_test.DidNavigateToPage(page, tab)
+
+  def WillRunActions(self, page, tab):
+    if self.page_test:
+      self.page_test.WillRunActions(page, tab)
+
+  def DidRunActions(self, page, tab):
+    if self.page_test:
+      self.page_test.DidRunActions(page, tab)
+
+  def ValidatePage(self, page, tab, results):
+    if self.page_test:
+      self.page_test.ValidatePage(page, tab, results)
 
   def RunPage(self, page, tab, results):
     tab.WaitForDocumentReadyStateToBeComplete()
@@ -58,96 +63,138 @@ class RecordPage(page_test.PageTest):  # pylint: disable=W0223
     # speed index metric.
     time.sleep(3)
 
-    # Run the actions for all measurements. Reload the page between
-    # actions.
+    # When running record_wpr, results is a GTestTestResults, so we create a
+    # dummy PageMeasurementResults that implements the functions we use.
+    # TODO(chrishenry): Fix the need for a dummy_results object.
+    dummy_results = page_measurement_results.PageMeasurementResults()
+
+    if self.page_test:
+      self._action_name_to_run = self.page_test.action_name_to_run
+      self.page_test.RunPage(page, tab, dummy_results)
+      return
+
     should_reload = False
-    interactive = self.options and self.options.interactive
+    # Run the actions on the page for all available measurements.
     for action_name in self._action_names:
+      # Skip this action if it is not defined
       if not hasattr(page, action_name):
         continue
+      # Reload the page between actions to start with a clean slate.
       if should_reload:
         self.RunNavigateSteps(page, tab)
-      action_runner = action_runner_module.ActionRunner(tab)
-      if interactive:
-        action_runner.PauseInteractive()
-      else:
-        self._RunMethod(page, action_name, action_runner)
+      self._action_name_to_run = action_name
+      super(RecorderPageTest, self).RunPage(page, tab, dummy_results)
       should_reload = True
 
-    # Run the PageTest's validator, so that we capture any additional resources
-    # that are loaded by the test.
-    if self.test:
-      dummy_results = page_measurement_results.PageMeasurementResults()
-      self.test.ValidatePage(page, tab, dummy_results)
+  def RunNavigateSteps(self, page, tab):
+    if self.page_test:
+      self.page_test.RunNavigateSteps(page, tab)
+    else:
+      super(RecorderPageTest, self).RunNavigateSteps(page, tab)
+
+
+def FindAllActionNames(base_dir):
+  """Returns a set of of all action names used in our measurements."""
+  action_names = set()
+  # Get all PageMeasurements except for ProfileCreators (see crbug.com/319573)
+  for _, cls in discover.DiscoverClasses(
+      base_dir, base_dir, page_measurement.PageMeasurement).items():
+    if not issubclass(cls, profile_creator.ProfileCreator):
+      action_name = cls().action_name_to_run
+      if action_name:
+        action_names.add(action_name)
+  return action_names
+
+
+def _MaybeGetInstanceOfClass(target, base_dir, cls):
+  if isinstance(target, cls):
+    return target
+  classes = discover.DiscoverClasses(base_dir, base_dir, cls,
+                                     index_by_class_name=True)
+  return classes[target]() if target in classes else None
+
+
+class WprRecorder(object):
+
+  def __init__(self, base_dir, target, extra_args=None):
+    action_names_to_run = FindAllActionNames(base_dir)
+    self._record_page_test = RecorderPageTest(action_names_to_run)
+    self._temp_target_wpr_file_path = tempfile.mkstemp()[1]
+    self._options = self._CreateOptions()
+
+    self._benchmark = _MaybeGetInstanceOfClass(target, base_dir,
+                                               benchmark.Benchmark)
+    if self._benchmark is not None:
+      self._record_page_test.page_test = self._benchmark.test()
+    self._parser = self._options.CreateParser(usage='%prog <PageSet|Benchmark>')
+    self._AddCommandLineArgs()
+    self._ParseArgs(extra_args)
+    self._ProcessCommandLineArgs()
+    self._page_set = self._GetPageSet(base_dir, target)
+
+  @property
+  def options(self):
+    return self._options
+
+  def _CreateOptions(self):
+    options = browser_options.BrowserFinderOptions()
+    options.browser_options.wpr_mode = wpr_modes.WPR_RECORD
+    options.browser_options.no_proxy_server = True
+    return options
+
+  def _AddCommandLineArgs(self):
+    page_runner.AddCommandLineArgs(self._parser)
+    if self._benchmark is not None:
+      self._benchmark.AddCommandLineArgs(self._parser)
+      self._benchmark.SetArgumentDefaults(self._parser)
+
+  def _ParseArgs(self, extra_args=None):
+    args = sys.argv[1:]
+    if extra_args is not None:
+      args += extra_args
+    self._parser.parse_args(args)
+
+  def _ProcessCommandLineArgs(self):
+    page_runner.ProcessCommandLineArgs(self._parser, self._options)
+    if self._benchmark is not None:
+      self._benchmark.ProcessCommandLineArgs(self._parser, self._options)
+
+  def _GetPageSet(self, base_dir, target):
+    if self._benchmark is not None:
+      return self._benchmark.CreatePageSet(self._options)
+    ps = _MaybeGetInstanceOfClass(target, base_dir, page_set.PageSet)
+    if ps is None:
+      self._parser.print_usage()
+      sys.exit(1)
+    return ps
+
+  def Record(self):
+    self._page_set.wpr_archive_info.AddNewTemporaryRecording(
+        self._temp_target_wpr_file_path)
+    self._record_page_test.CustomizeBrowserOptions(self._options)
+    return page_runner.Run(self._record_page_test, self._page_set,
+                           test_expectations.TestExpectations(), self._options)
+
+  def HandleResults(self, results):
+    if results.failures or results.skipped:
+      logging.warning('Some pages failed and/or were skipped. The recording '
+                      'has not been updated for these pages.')
+    results.PrintSummary()
+
+    if results.successes:
+      # Update the metadata for the pages which were recorded.
+      self._page_set.wpr_archive_info.AddRecordedPages(results.successes)
+    else:
+      os.remove(self._temp_target_wpr_file_path)
 
 
 def Main(base_dir):
-  measurements = {
-      n: cls for n, cls in discover.DiscoverClasses(
-          base_dir, base_dir, page_measurement.PageMeasurement).items()
-      # Filter out unneeded ProfileCreators (crbug.com/319573).
-      if not issubclass(cls, profile_creator.ProfileCreator)
-      }
-  tests = discover.DiscoverClasses(base_dir, base_dir, benchmark.Benchmark,
-                                   index_by_class_name=True)
-
-  options = browser_options.BrowserFinderOptions()
-  parser = options.CreateParser('%prog <PageSet|Test|URL>')
-  page_runner.AddCommandLineArgs(parser)
-
-  recorder = RecordPage(measurements)
-  recorder.AddCommandLineArgs(parser)
-
   quick_args = [a for a in sys.argv[1:] if not a.startswith('-')]
   if len(quick_args) != 1:
-    parser.print_usage()
+    print >> sys.stderr, 'Usage: record_wpr <PageSet|Benchmark>\n'
     sys.exit(1)
-  target = quick_args[0]
-  if target in tests:
-    recorder.test = tests[target]().test()
-    recorder.test.AddCommandLineArgs(parser)
-    recorder.test.SetArgumentDefaults(parser)
-    parser.parse_args()
-    recorder.test.ProcessCommandLineArgs(parser, options)
-    ps = tests[target]().CreatePageSet(options)
-  elif discover.IsPageSetFile(target):
-    parser.parse_args()
-    ps = page_set.PageSet.FromFile(target)
-  else:
-    parser.print_usage()
-    sys.exit(1)
-
-  page_runner.ProcessCommandLineArgs(parser, options)
-  recorder.ProcessCommandLineArgs(parser, options)
-
-  expectations = test_expectations.TestExpectations()
-
-  # Set the archive path to something temporary.
-  temp_target_wpr_file_path = tempfile.mkstemp()[1]
-  ps.wpr_archive_info.AddNewTemporaryRecording(temp_target_wpr_file_path)
-
-  # Do the actual recording.
-  options.browser_options.wpr_mode = wpr_modes.WPR_RECORD
-  options.browser_options.no_proxy_server = True
-  recorder.CustomizeBrowserOptions(options)
-  results = page_runner.Run(recorder, ps, expectations, options)
-
-  if results.failures:
-    logging.warning('Some pages failed. The recording has not been updated for '
-                    'these pages.')
-    logging.warning('Failed pages:\n%s', '\n'.join(
-        p.display_name for p in results.pages_that_had_failures))
-
-  if results.skipped:
-    logging.warning('Some pages were skipped. The recording has not been '
-                    'updated for these pages.')
-    logging.warning('Skipped pages:\n%s', '\n'.join(
-        p.display_name for p in zip(*results.skipped)[0]))
-
-  if results.successes:
-    # Update the metadata for the pages which were recorded.
-    ps.wpr_archive_info.AddRecordedPages(results.successes)
-  else:
-    os.remove(temp_target_wpr_file_path)
-
+  target = quick_args.pop()
+  wpr_recorder = WprRecorder(base_dir, target)
+  results = wpr_recorder.Record()
+  wpr_recorder.HandleResults(results)
   return min(255, len(results.failures))