Refactor of record_wpr.py

* Breaks major steps out into functions
* Adds additional tests
* Adds support for command line flags defined in benchmarks
* Calls the PageTest RunPage directly, rather than kinda sorta reimplementing parts ourselves
* Fixes running benchmarks with record_wpr.py (crbug.com/378064)
* Page sets are referenced via class name on the command line, rather than by filename.

Other minor changes:
* Output will be different for failed pages at the end, using the default results PrintSummary() rather than a separate implementation (in this case GTestTestResults).

BUG=378064,367292

Review URL: https://codereview.chromium.org/374793002

git-svn-id: svn://svn.chromium.org/chrome/trunk/src@285490 0039d316-1c4b-4281-b951-d872f2087c98

6 files changed, 355 insertions, 133 deletions

diff --git a/tools/telemetry/telemetry/page/record_wpr.py b/tools/telemetry/telemetry/page/record_wpr.py
index 9491e4d..561a7f9 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))
diff --git a/tools/telemetry/telemetry/page/record_wpr_unittest.py b/tools/telemetry/telemetry/page/record_wpr_unittest.py
index 4c89143..90b54c3 100644
--- a/tools/telemetry/telemetry/page/record_wpr_unittest.py
+++ b/tools/telemetry/telemetry/page/record_wpr_unittest.py
@@ -2,56 +2,189 @@
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
-import unittest
+import os
+import sys
 
+from telemetry import benchmark
+from telemetry.core import util
+from telemetry.core import wpr_modes
 from telemetry.page import page as page_module
+from telemetry.page import page_set as page_set_module
+from telemetry.page import page_test
 from telemetry.page import record_wpr
+from telemetry.unittest import tab_test_case
 
 
-class TestPage(page_module.Page):
-  def __init__(self):
-    super(TestPage, self).__init__(url='file://foo.html',
-                                   page_set=None,
-                                   base_dir=None)
-    self.run_navigate = False
-    self.run_foo = False
-    self.run_bar = False
+class MockPage(page_module.Page):
+  def __init__(self, page_set, url):
+    super(MockPage, self).__init__(url=url,
+                                   page_set=page_set,
+                                   base_dir=util.GetUnittestDataDir())
+    self.func_calls = []
 
-  def RunNavigateSteps(self, _):
-    self.run_navigate = True
+  def RunNavigateSteps(self, action_runner):
+    self.func_calls.append('RunNavigateSteps')
+    super(MockPage, self).RunNavigateSteps(action_runner)
 
   def RunFoo(self, _):
-    self.run_foo = True
+    self.func_calls.append('RunFoo')
 
   def RunBar(self, _):
-    self.run_bar = True
-
+    self.func_calls.append('RunBar')
 
-class FakeFooMeasurement(object):
-  def __init__(self):
-    self.action_name_to_run = "RunFoo"
+  def RunBaz(self, _):
+    self.func_calls.append('RunBaz')
 
 
-class FakeBarMeasurement(object):
-  def __init__(self):
-    self.action_name_to_run = "RunBar"
-
+class MockPageSet(page_set_module.PageSet):
+  def __init__(self, url=''):
+    super(MockPageSet, self).__init__(archive_data_file='data/test.json')
+    self.AddPage(MockPage(self, url))
 
-class FakeTab(object):
-  def WaitForDocumentReadyStateToBeComplete(self):
-    pass
-
-
-class RecordWprUnitTest(unittest.TestCase):
-  def setUp(self):
-    super(RecordWprUnitTest, self).setUp()
-
-  def testRunActions(self):
-    page = TestPage()
-    record_runner = record_wpr.RecordPage({1 : FakeFooMeasurement,
-                                           2 : FakeBarMeasurement})
-    record_runner.RunPage(page, tab=FakeTab(), results=None)
-    self.assertTrue(page.run_navigate)
-    self.assertTrue(page.run_foo)
-    self.assertTrue(page.run_bar)
 
+class MockPageTest(page_test.PageTest):
+  def __init__(self):
+    super(MockPageTest, self).__init__()
+    self._action_name_to_run = "RunBaz"
+    self.func_calls = []
+
+  @classmethod
+  def AddCommandLineArgs(cls, parser):
+    parser.add_option('--mock-page-test-option', action="store_true")
+
+  def WillNavigateToPage(self, page, tab):
+    self.func_calls.append('WillNavigateToPage')
+
+  def DidNavigateToPage(self, page, tab):
+    self.func_calls.append('DidNavigateToPage')
+
+  def WillRunActions(self, page, tab):
+    self.func_calls.append('WillRunActions')
+
+  def DidRunActions(self, page, tab):
+    self.func_calls.append('DidRunActions')
+
+  def ValidatePage(self, page, tab, results):
+    self.func_calls.append('ValidatePage')
+
+
+class MockBenchmark(benchmark.Benchmark):
+  test = MockPageTest
+
+  @classmethod
+  def AddTestCommandLineArgs(cls, group):
+    group.add_option('', '--mock-benchmark-url', action='store', type='string')
+
+  def CreatePageSet(self, options):
+    kwargs = {}
+    if (options.mock_benchmark_url):
+      kwargs['url'] = options.mock_benchmark_url
+    return MockPageSet(**kwargs)
+
+
+class RecordWprUnitTests(tab_test_case.TabTestCase):
+
+  _base_dir = util.GetUnittestDataDir()
+  _test_data_dir = os.path.join(util.GetUnittestDataDir(), 'page_measurements')
+
+  @classmethod
+  def setUpClass(cls):
+    sys.path.extend([cls._base_dir, cls._test_data_dir])
+    super(RecordWprUnitTests, cls).setUpClass()
+    cls._browser.SetHTTPServerDirectories(util.GetUnittestDataDir())
+    blank_html_path = os.path.join(util.GetUnittestDataDir(), 'blank.html')
+    cls._url = cls._browser.http_server.UrlOf(blank_html_path)
+
+  # When the RecorderPageTest is created from a PageSet, we do not have a
+  # PageTest to use. In this case, we will record every available action.
+  def testRunPage_AllActions(self):
+    record_page_test = record_wpr.RecorderPageTest(["RunFoo", "RunBar"])
+    page = MockPage(page_set=MockPageSet(url=self._url), url=self._url)
+    record_page_test.RunPage(page, self._tab, results=None)
+    self.assertTrue('RunFoo' in page.func_calls)
+    self.assertTrue('RunBar' in page.func_calls)
+    self.assertFalse('RunBaz' in page.func_calls)
+
+  def testRunPage_DontReloadSingleActions(self):
+    record_page_test = record_wpr.RecorderPageTest(["RunFoo"])
+    page = MockPage(page_set=MockPageSet(url=self._url), url=self._url)
+    record_page_test.RunPage(page, self._tab, results=None)
+    self.assertFalse('RunNavigateSteps' in page.func_calls)
+
+  def testRunPage_ReloadPageBetweenActions(self):
+    record_page_test = record_wpr.RecorderPageTest(["RunFoo", "RunBar"])
+    page = MockPage(page_set=MockPageSet(url=self._url), url=self._url)
+    record_page_test.RunPage(page, self._tab, results=None)
+    self.assertTrue('RunNavigateSteps' in page.func_calls)
+
+  # When the RecorderPageTest is created from a Benchmark, the benchmark will
+  # have a PageTest, specified by its test attribute.
+  def testRunPage_OnlyRunBenchmarkAction(self):
+    record_page_test = record_wpr.RecorderPageTest(["RunFoo"])
+    record_page_test.page_test = MockBenchmark().test()
+    page = MockPage(page_set=MockPageSet(url=self._url), url=self._url)
+    record_page_test.RunPage(page, self._tab, results=None)
+    self.assertFalse('RunFoo' in page.func_calls)
+    self.assertTrue('RunBaz' in page.func_calls)
+
+  def testRunPage_CallBenchmarksPageTestsFunctions(self):
+    record_page_test = record_wpr.RecorderPageTest([])
+    record_page_test.page_test = MockBenchmark().test()
+    page = MockPage(page_set=MockPageSet(url=self._url), url=self._url)
+    record_page_test.RunPage(page, self._tab, results=None)
+    self.assertEqual(3, len(record_page_test.page_test.func_calls))
+    self.assertEqual('WillRunActions', record_page_test.page_test.func_calls[0])
+    self.assertEqual('DidRunActions', record_page_test.page_test.func_calls[1])
+    self.assertEqual('ValidatePage', record_page_test.page_test.func_calls[2])
+
+  def testWprRecorderWithPageSet(self):
+    wpr_recorder = record_wpr.WprRecorder(self._test_data_dir,
+                                          MockPageSet(url=self._url))
+    results = wpr_recorder.Record()
+    self.assertEquals(1, len(results.successes))
+    mock_page = results.successes.pop()
+    self.assertTrue('RunFoo' in mock_page.func_calls)
+    self.assertFalse('RunBaz' in mock_page.func_calls)
+
+  def testWprRecorderWithBenchmark(self):
+    flags = ['--mock-benchmark-url', self._url]
+    wpr_recorder = record_wpr.WprRecorder(self._test_data_dir, MockBenchmark(),
+                                          flags)
+    results = wpr_recorder.Record()
+    self.assertEquals(1, len(results.successes))
+    mock_page = results.successes.pop()
+    self.assertFalse('RunFoo' in mock_page.func_calls)
+    self.assertTrue('RunBaz' in mock_page.func_calls)
+
+  def testCommandLineFlags(self):
+    flags = [
+        '--page-repeat', '2',
+        '--mock-benchmark-url', self._url,
+        '--mock-page-test-option',
+    ]
+    wpr_recorder = record_wpr.WprRecorder(self._test_data_dir, MockBenchmark(),
+                                          flags)
+    # page_runner command-line args
+    self.assertEquals(2, wpr_recorder.options.page_repeat)
+    # benchmark command-line args
+    self.assertEquals(self._url, wpr_recorder.options.mock_benchmark_url)
+    # benchmark's page_test command-line args
+    self.assertTrue(wpr_recorder.options.mock_page_test_option)
+    # invalid command-line args
+    self.assertFalse(hasattr(wpr_recorder.options, 'not_a_real_option'))
+
+  def testRecordingEnabled(self):
+    flags = ['--mock-benchmark-url', self._url]
+    wpr_recorder = record_wpr.WprRecorder(self._test_data_dir, MockBenchmark(),
+                                          flags)
+    self.assertEqual(wpr_modes.WPR_RECORD,
+                     wpr_recorder.options.browser_options.wpr_mode)
+
+  def testFindAllActionNames(self):
+    # The src/tools/telemetry/unittest_data/page_measurements/ has been
+    # populated with three simple Page Measurement classes, the first two of
+    # which have action_name_to_run defined.
+    action_names_to_run = record_wpr.FindAllActionNames(self._test_data_dir)
+    self.assertTrue('RunFoo' in action_names_to_run)
+    self.assertTrue('RunBar' in action_names_to_run)
+    self.assertFalse('RunBaz' in action_names_to_run)
diff --git a/tools/telemetry/unittest_data/page_measurements/__init__.py b/tools/telemetry/unittest_data/page_measurements/__init__.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/tools/telemetry/unittest_data/page_measurements/__init__.py
diff --git a/tools/telemetry/unittest_data/page_measurements/page_measurement1.py b/tools/telemetry/unittest_data/page_measurements/page_measurement1.py
new file mode 100644
index 0000000..dd1068e
--- /dev/null
+++ b/tools/telemetry/unittest_data/page_measurements/page_measurement1.py
@@ -0,0 +1,14 @@
+# Copyright 2014 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.
+
+"""A simple PageMeasurement used by page/record_wpr.py's unit tests."""
+
+from telemetry.page import page_measurement
+
+class MockPageMeasurementOne(page_measurement.PageMeasurement):
+  def __init__(self):
+    super(MockPageMeasurementOne, self).__init__(action_name_to_run="RunFoo")
+
+  def MeasurePage(self, page, tab, results):
+    pass
diff --git a/tools/telemetry/unittest_data/page_measurements/page_measurement2.py b/tools/telemetry/unittest_data/page_measurements/page_measurement2.py
new file mode 100644
index 0000000..3f13e69
--- /dev/null
+++ b/tools/telemetry/unittest_data/page_measurements/page_measurement2.py
@@ -0,0 +1,14 @@
+# Copyright 2014 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.
+
+"""A simple PageMeasurement used by page/record_wpr.py's unit tests."""
+
+from telemetry.page import page_measurement
+
+class MockPageMeasurementTwo(page_measurement.PageMeasurement):
+  def __init__(self):
+    super(MockPageMeasurementTwo, self).__init__(action_name_to_run="RunBar")
+
+  def MeasurePage(self, page, tab, results):
+    pass
diff --git a/tools/telemetry/unittest_data/page_measurements/page_measurement3.py b/tools/telemetry/unittest_data/page_measurements/page_measurement3.py
new file mode 100644
index 0000000..083c91b
--- /dev/null
+++ b/tools/telemetry/unittest_data/page_measurements/page_measurement3.py
@@ -0,0 +1,14 @@
+# Copyright 2014 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.
+
+"""A simple PageMeasurement used by page/record_wpr.py's unit tests."""
+
+from telemetry.page import page_measurement
+
+class MockPageMeasurementThree(page_measurement.PageMeasurement):
+  def __init__(self):
+    super(MockPageMeasurementThree, self).__init__(action_name_to_run=None)
+
+  def MeasurePage(self, page, tab, results):
+    pass