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 argparse
6import logging
7import sys
8
9from telemetry import benchmark
10from telemetry import story
11from telemetry.core import discover
12from telemetry.core import util
13from telemetry.internal.browser import browser_options
14from telemetry.internal.results import results_options
15from telemetry.internal import story_runner
16from telemetry.internal.util import binary_manager
17from telemetry.page import page_test
18from telemetry.util import matching
19from telemetry.util import wpr_modes
20from telemetry.web_perf import timeline_based_measurement
21from telemetry.web_perf import timeline_based_page_test
22
23
24class RecorderPageTest(page_test.PageTest):
25  def __init__(self):
26    super(RecorderPageTest, self).__init__()
27    self.page_test = None
28
29  def CustomizeBrowserOptions(self, options):
30    if self.page_test:
31      self.page_test.CustomizeBrowserOptions(options)
32
33  def WillStartBrowser(self, browser):
34    if self.page_test:
35      self.page_test.WillStartBrowser(browser)
36
37  def DidStartBrowser(self, browser):
38    if self.page_test:
39      self.page_test.DidStartBrowser(browser)
40
41  def WillNavigateToPage(self, page, tab):
42    """Override to ensure all resources are fetched from network."""
43    tab.ClearCache(force=False)
44    if self.page_test:
45      self.page_test.WillNavigateToPage(page, tab)
46
47  def DidNavigateToPage(self, page, tab):
48    if self.page_test:
49      self.page_test.DidNavigateToPage(page, tab)
50    tab.WaitForDocumentReadyStateToBeComplete()
51    util.WaitFor(tab.HasReachedQuiescence, 30)
52
53  def CleanUpAfterPage(self, page, tab):
54    if self.page_test:
55      self.page_test.CleanUpAfterPage(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 RunNavigateSteps(self, page, tab):
62    if self.page_test:
63      self.page_test.RunNavigateSteps(page, tab)
64    else:
65      super(RecorderPageTest, self).RunNavigateSteps(page, tab)
66
67
68def _GetSubclasses(base_dir, cls):
69  """Returns all subclasses of |cls| in |base_dir|.
70
71  Args:
72    cls: a class
73
74  Returns:
75    dict of {underscored_class_name: benchmark class}
76  """
77  return discover.DiscoverClasses(base_dir, base_dir, cls,
78                                  index_by_class_name=True)
79
80
81def _MaybeGetInstanceOfClass(target, base_dir, cls):
82  if isinstance(target, cls):
83    return target
84  classes = _GetSubclasses(base_dir, cls)
85  return classes[target]() if target in classes else None
86
87
88def _PrintAllImpl(all_items, item_name, output_stream):
89  output_stream.write('Available %s\' names with descriptions:\n' % item_name)
90  keys = sorted(all_items.keys())
91  key_description = [(k, all_items[k].Description()) for k in keys]
92  _PrintPairs(key_description, output_stream)
93  output_stream.write('\n')
94
95
96def _PrintAllBenchmarks(base_dir, output_stream):
97  # TODO: reuse the logic of finding supported benchmarks in benchmark_runner.py
98  # so this only prints out benchmarks that are supported by the recording
99  # platform.
100  _PrintAllImpl(_GetSubclasses(base_dir, benchmark.Benchmark), 'benchmarks',
101                output_stream)
102
103
104def _PrintAllStories(base_dir, output_stream):
105  # TODO: actually print all stories once record_wpr support general
106  # stories recording.
107  _PrintAllImpl(_GetSubclasses(base_dir, story.StorySet), 'story sets',
108                output_stream)
109
110
111def _PrintPairs(pairs, output_stream, prefix=''):
112  """Prints a list of string pairs with alignment."""
113  first_column_length = max(len(a) for a, _ in pairs)
114  format_string = '%s%%-%ds  %%s\n' % (prefix, first_column_length)
115  for a, b in pairs:
116    output_stream.write(format_string % (a, b.strip()))
117
118
119class WprRecorder(object):
120
121  def __init__(self, base_dir, target, args=None):
122    self._base_dir = base_dir
123    self._record_page_test = RecorderPageTest()
124    self._options = self._CreateOptions()
125
126    self._benchmark = _MaybeGetInstanceOfClass(target, base_dir,
127                                               benchmark.Benchmark)
128    self._parser = self._options.CreateParser(usage='See %prog --help')
129    self._AddCommandLineArgs()
130    self._ParseArgs(args)
131    self._ProcessCommandLineArgs()
132    if self._benchmark is not None:
133      test = self._benchmark.CreatePageTest(self.options)
134      if isinstance(test, timeline_based_measurement.TimelineBasedMeasurement):
135        test = timeline_based_page_test.TimelineBasedPageTest(test)
136      # This must be called after the command line args are added.
137      self._record_page_test.page_test = test
138
139    self._page_set_base_dir = (
140        self._options.page_set_base_dir if self._options.page_set_base_dir
141        else self._base_dir)
142    self._story_set = self._GetStorySet(target)
143
144  @property
145  def options(self):
146    return self._options
147
148  def _CreateOptions(self):
149    options = browser_options.BrowserFinderOptions()
150    options.browser_options.wpr_mode = wpr_modes.WPR_RECORD
151    options.browser_options.no_proxy_server = True
152    return options
153
154  def CreateResults(self):
155    if self._benchmark is not None:
156      benchmark_metadata = self._benchmark.GetMetadata()
157    else:
158      benchmark_metadata = benchmark.BenchmarkMetadata('record_wpr')
159
160    return results_options.CreateResults(benchmark_metadata, self._options)
161
162  def _AddCommandLineArgs(self):
163    self._parser.add_option('--page-set-base-dir', action='store',
164                            type='string')
165    story_runner.AddCommandLineArgs(self._parser)
166    if self._benchmark is not None:
167      self._benchmark.AddCommandLineArgs(self._parser)
168      self._benchmark.SetArgumentDefaults(self._parser)
169    self._parser.add_option('--upload', action='store_true')
170    self._SetArgumentDefaults()
171
172  def _SetArgumentDefaults(self):
173    self._parser.set_defaults(**{'output_formats': ['none']})
174
175  def _ParseArgs(self, args=None):
176    args_to_parse = sys.argv[1:] if args is None else args
177    self._parser.parse_args(args_to_parse)
178
179  def _ProcessCommandLineArgs(self):
180    story_runner.ProcessCommandLineArgs(self._parser, self._options)
181
182    if self._options.use_live_sites:
183      self._parser.error("Can't --use-live-sites while recording")
184
185    if self._benchmark is not None:
186      self._benchmark.ProcessCommandLineArgs(self._parser, self._options)
187
188  def _GetStorySet(self, target):
189    if self._benchmark is not None:
190      return self._benchmark.CreateStorySet(self._options)
191    story_set = _MaybeGetInstanceOfClass(target, self._page_set_base_dir,
192                                         story.StorySet)
193    if story_set is None:
194      sys.stderr.write('Target %s is neither benchmark nor story set.\n'
195                       % target)
196      if not self._HintMostLikelyBenchmarksStories(target):
197        sys.stderr.write(
198            'Found no similar benchmark or story. Please use '
199            '--list-benchmarks or --list-stories to list candidates.\n')
200        self._parser.print_usage()
201      sys.exit(1)
202    return story_set
203
204  def _HintMostLikelyBenchmarksStories(self, target):
205    def _Impl(all_items, category_name):
206      candidates = matching.GetMostLikelyMatchedObject(
207          all_items.iteritems(), target, name_func=lambda kv: kv[1].Name())
208      if candidates:
209        sys.stderr.write('\nDo you mean any of those %s below?\n' %
210                         category_name)
211        _PrintPairs([(k, v.Description()) for k, v in candidates], sys.stderr)
212        return True
213      return False
214
215    has_benchmark_hint = _Impl(
216        _GetSubclasses(self._base_dir, benchmark.Benchmark), 'benchmarks')
217    has_story_hint = _Impl(
218        _GetSubclasses(self._base_dir, story.StorySet), 'stories')
219    return has_benchmark_hint or has_story_hint
220
221  def Record(self, results):
222    assert self._story_set.wpr_archive_info, (
223      'Pageset archive_data_file path must be specified.')
224    self._story_set.wpr_archive_info.AddNewTemporaryRecording()
225    self._record_page_test.CustomizeBrowserOptions(self._options)
226    story_runner.Run(self._record_page_test, self._story_set,
227        self._options, results)
228
229  def HandleResults(self, results, upload_to_cloud_storage):
230    if results.failures or results.skipped_values:
231      logging.warning('Some pages failed and/or were skipped. The recording '
232                      'has not been updated for these pages.')
233    results.PrintSummary()
234    self._story_set.wpr_archive_info.AddRecordedStories(
235        results.pages_that_succeeded,
236        upload_to_cloud_storage)
237
238
239def Main(environment):
240
241  parser = argparse.ArgumentParser(
242      usage='Record a benchmark or a story (page set).')
243  parser.add_argument(
244      'benchmark',
245      help=('benchmark name. This argument is optional. If both benchmark name '
246            'and story name are specified, this takes precedence as the '
247            'target of the recording.'),
248      nargs='?')
249  parser.add_argument('--story', help='story (page set) name')
250  parser.add_argument('--list-stories', dest='list_stories',
251                      action='store_true', help='list all story names.')
252  parser.add_argument('--list-benchmarks', dest='list_benchmarks',
253                      action='store_true', help='list all benchmark names.')
254  parser.add_argument('--upload', action='store_true',
255                      help='upload to cloud storage.')
256  args, extra_args = parser.parse_known_args()
257
258  if args.list_benchmarks or args.list_stories:
259    if args.list_benchmarks:
260      _PrintAllBenchmarks(environment.top_level_dir, sys.stderr)
261    if args.list_stories:
262      _PrintAllStories(environment.top_level_dir, sys.stderr)
263    return 0
264
265  target = args.benchmark or args.story
266
267  if not target:
268    sys.stderr.write('Please specify target (benchmark or story). Please refer '
269                     'usage below\n\n')
270    parser.print_help()
271    return 0
272
273  binary_manager.InitDependencyManager(environment.client_config)
274
275
276  # TODO(nednguyen): update WprRecorder so that it handles the difference
277  # between recording a benchmark vs recording a story better based on
278  # the distinction between args.benchmark & args.story
279  wpr_recorder = WprRecorder(environment.top_level_dir, target, extra_args)
280  results = wpr_recorder.CreateResults()
281  wpr_recorder.Record(results)
282  wpr_recorder.HandleResults(results, args.upload)
283  return min(255, len(results.failures))
284