1# Copyright 2012 The Chromium Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5import logging
6import sys
7
8from telemetry import benchmark
9from telemetry.core import browser_options
10from telemetry.core import discover
11from telemetry.core import util
12from telemetry.core import wpr_modes
13from telemetry.page import page_runner
14from telemetry.page import page_set
15from telemetry.page import page_test
16from telemetry.page import profile_creator
17from telemetry.page import test_expectations
18from telemetry.results import results_options
19
20
21class RecorderPageTest(page_test.PageTest):  # pylint: disable=W0223
22  def __init__(self, action_names):
23    super(RecorderPageTest, self).__init__()
24    self._action_names = action_names
25    self.page_test = None
26
27  def CanRunForPage(self, page):
28    return page.url.startswith('http')
29
30  def WillStartBrowser(self, browser):
31    if self.page_test:
32      self.page_test.WillStartBrowser(browser)
33
34  def DidStartBrowser(self, browser):
35    if self.page_test:
36      self.page_test.DidStartBrowser(browser)
37
38  def WillNavigateToPage(self, page, tab):
39    """Override to ensure all resources are fetched from network."""
40    tab.ClearCache(force=False)
41    if self.page_test:
42      self.page_test.options = self.options
43      self.page_test.WillNavigateToPage(page, tab)
44
45  def DidNavigateToPage(self, page, tab):
46    if self.page_test:
47      self.page_test.DidNavigateToPage(page, tab)
48
49  def WillRunActions(self, page, tab):
50    if self.page_test:
51      self.page_test.WillRunActions(page, tab)
52
53  def DidRunActions(self, page, tab):
54    if self.page_test:
55      self.page_test.DidRunActions(page, tab)
56
57  def ValidateAndMeasurePage(self, page, tab, results):
58    if self.page_test:
59      self.page_test.ValidateAndMeasurePage(page, tab, results)
60
61  def RunPage(self, page, tab, results):
62    tab.WaitForDocumentReadyStateToBeComplete()
63    util.WaitFor(tab.HasReachedQuiescence, 30)
64
65    if self.page_test:
66      self._action_name_to_run = self.page_test.action_name_to_run
67      self.page_test.RunPage(page, tab, results)
68      return
69
70    should_reload = False
71    # Run the actions on the page for all available measurements.
72    for action_name in self._action_names:
73      # Skip this action if it is not defined
74      if not hasattr(page, action_name):
75        continue
76      # Reload the page between actions to start with a clean slate.
77      if should_reload:
78        self.RunNavigateSteps(page, tab)
79      self._action_name_to_run = action_name
80      super(RecorderPageTest, self).RunPage(page, tab, results)
81      should_reload = True
82
83  def RunNavigateSteps(self, page, tab):
84    if self.page_test:
85      self.page_test.RunNavigateSteps(page, tab)
86    else:
87      super(RecorderPageTest, self).RunNavigateSteps(page, tab)
88
89
90def FindAllActionNames(base_dir):
91  """Returns a set of of all action names used in our measurements."""
92  action_names = set()
93  # Get all PageTests except for ProfileCreators (see crbug.com/319573)
94  for _, cls in discover.DiscoverClasses(
95      base_dir, base_dir, page_test.PageTest).items():
96    if not issubclass(cls, profile_creator.ProfileCreator):
97      action_name = cls().action_name_to_run
98      if action_name:
99        action_names.add(action_name)
100  return action_names
101
102
103def _MaybeGetInstanceOfClass(target, base_dir, cls):
104  if isinstance(target, cls):
105    return target
106  classes = discover.DiscoverClasses(base_dir, base_dir, cls,
107                                     index_by_class_name=True)
108  return classes[target]() if target in classes else None
109
110
111class WprRecorder(object):
112
113  def __init__(self, base_dir, target, args=None):
114    action_names_to_run = FindAllActionNames(base_dir)
115    self._record_page_test = RecorderPageTest(action_names_to_run)
116    self._options = self._CreateOptions()
117
118    self._benchmark = _MaybeGetInstanceOfClass(target, base_dir,
119                                               benchmark.Benchmark)
120    if self._benchmark is not None:
121      self._record_page_test.page_test = self._benchmark.test()
122    self._parser = self._options.CreateParser(usage='%prog <PageSet|Benchmark>')
123    self._AddCommandLineArgs()
124    self._ParseArgs(args)
125    self._ProcessCommandLineArgs()
126    self._page_set = self._GetPageSet(base_dir, target)
127
128  @property
129  def options(self):
130    return self._options
131
132  def _CreateOptions(self):
133    options = browser_options.BrowserFinderOptions()
134    options.browser_options.wpr_mode = wpr_modes.WPR_RECORD
135    options.browser_options.no_proxy_server = True
136    return options
137
138  def CreateResults(self):
139    if self._benchmark is not None:
140      benchmark_metadata = self._benchmark.GetMetadata()
141    else:
142      benchmark_metadata = benchmark.BenchmarkMetadata('record_wpr')
143
144    return results_options.CreateResults(benchmark_metadata, self._options)
145
146  def _AddCommandLineArgs(self):
147    page_runner.AddCommandLineArgs(self._parser)
148    if self._benchmark is not None:
149      self._benchmark.AddCommandLineArgs(self._parser)
150      self._benchmark.SetArgumentDefaults(self._parser)
151    self._SetArgumentDefaults()
152
153  def _SetArgumentDefaults(self):
154    self._parser.set_defaults(**{'output_formats': ['none']})
155
156  def _ParseArgs(self, args=None):
157    args_to_parse = sys.argv[1:] if args is None else args
158    self._parser.parse_args(args_to_parse)
159
160  def _ProcessCommandLineArgs(self):
161    page_runner.ProcessCommandLineArgs(self._parser, self._options)
162    if self._benchmark is not None:
163      self._benchmark.ProcessCommandLineArgs(self._parser, self._options)
164
165  def _GetPageSet(self, base_dir, target):
166    if self._benchmark is not None:
167      return self._benchmark.CreatePageSet(self._options)
168    ps = _MaybeGetInstanceOfClass(target, base_dir, page_set.PageSet)
169    if ps is None:
170      self._parser.print_usage()
171      sys.exit(1)
172    return ps
173
174  def Record(self, results):
175    self._page_set.wpr_archive_info.AddNewTemporaryRecording()
176    self._record_page_test.CustomizeBrowserOptions(self._options)
177    page_runner.Run(self._record_page_test, self._page_set,
178        test_expectations.TestExpectations(), self._options, results)
179
180  def HandleResults(self, results):
181    if results.failures or results.skipped_values:
182      logging.warning('Some pages failed and/or were skipped. The recording '
183                      'has not been updated for these pages.')
184    results.PrintSummary()
185    self._page_set.wpr_archive_info.AddRecordedPages(
186        results.pages_that_succeeded)
187
188
189def Main(base_dir):
190  quick_args = [a for a in sys.argv[1:] if not a.startswith('-')]
191  if len(quick_args) != 1:
192    print >> sys.stderr, 'Usage: record_wpr <PageSet|Benchmark>\n'
193    sys.exit(1)
194  target = quick_args.pop()
195  wpr_recorder = WprRecorder(base_dir, target)
196  results = wpr_recorder.CreateResults()
197  wpr_recorder.Record(results)
198  wpr_recorder.HandleResults(results)
199  return min(255, len(results.failures))
200