1# Copyright 2014 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 optparse
7import os
8import shutil
9import sys
10import zipfile
11
12from telemetry import decorators
13from telemetry.core import browser_finder
14from telemetry.core import command_line
15from telemetry.core import util
16from telemetry.page import page_runner
17from telemetry.page import page_set
18from telemetry.page import page_test
19from telemetry.page import test_expectations
20from telemetry.results import results_options
21from telemetry.util import cloud_storage
22
23Disabled = decorators.Disabled
24Enabled = decorators.Enabled
25
26
27class BenchmarkMetadata(object):
28  def __init__(self, name):
29    self._name = name
30
31  @property
32  def name(self):
33    return self._name
34
35class Benchmark(command_line.Command):
36  """Base class for a Telemetry benchmark.
37
38  A test packages a PageTest and a PageSet together.
39  """
40  options = {}
41
42  @classmethod
43  def Name(cls):
44    name = cls.__module__.split('.')[-1]
45    if hasattr(cls, 'tag'):
46      name += '.' + cls.tag
47    if hasattr(cls, 'page_set'):
48      name += '.' + cls.page_set.Name()
49    return name
50
51  @classmethod
52  def AddCommandLineArgs(cls, parser):
53    cls.PageTestClass().AddCommandLineArgs(parser)
54
55    if hasattr(cls, 'AddTestCommandLineArgs'):
56      group = optparse.OptionGroup(parser, '%s test options' % cls.Name())
57      cls.AddTestCommandLineArgs(group)
58      parser.add_option_group(group)
59
60  @classmethod
61  def SetArgumentDefaults(cls, parser):
62    cls.PageTestClass().SetArgumentDefaults(parser)
63    parser.set_defaults(**cls.options)
64
65  @classmethod
66  def ProcessCommandLineArgs(cls, parser, args):
67    cls.PageTestClass().ProcessCommandLineArgs(parser, args)
68
69  def CustomizeBrowserOptions(self, options):
70    """Add browser options that are required by this benchmark."""
71
72  def GetMetadata(self):
73    return BenchmarkMetadata(self.Name())
74
75  def Run(self, finder_options):
76    """Run this test with the given options."""
77    self.CustomizeBrowserOptions(finder_options.browser_options)
78
79    pt = self.PageTestClass()()
80    pt.__name__ = self.__class__.__name__
81
82    if hasattr(self, '_disabled_strings'):
83      pt._disabled_strings = self._disabled_strings
84    if hasattr(self, '_enabled_strings'):
85      pt._enabled_strings = self._enabled_strings
86
87    ps = self.CreatePageSet(finder_options)
88    expectations = self.CreateExpectations(ps)
89
90    self._DownloadGeneratedProfileArchive(finder_options)
91
92    benchmark_metadata = self.GetMetadata()
93    results = results_options.CreateResults(benchmark_metadata, finder_options)
94    try:
95      page_runner.Run(pt, ps, expectations, finder_options, results)
96    except page_test.TestNotSupportedOnPlatformFailure as failure:
97      logging.warning(str(failure))
98
99    results.PrintSummary()
100    return len(results.failures)
101
102  def _DownloadGeneratedProfileArchive(self, options):
103    """Download and extract profile directory archive if one exists."""
104    archive_name = getattr(self, 'generated_profile_archive', None)
105
106    # If attribute not specified, nothing to do.
107    if not archive_name:
108      return
109
110    # If profile dir specified on command line, nothing to do.
111    if options.browser_options.profile_dir:
112      logging.warning("Profile directory specified on command line: %s, this"
113          "overrides the benchmark's default profile directory.",
114          options.browser_options.profile_dir)
115      return
116
117    # Download profile directory from cloud storage.
118    found_browser = browser_finder.FindBrowser(options)
119    test_data_dir = os.path.join(util.GetChromiumSrcDir(), 'tools', 'perf',
120        'generated_profiles',
121        found_browser.target_os)
122    generated_profile_archive_path = os.path.normpath(
123        os.path.join(test_data_dir, archive_name))
124
125    try:
126      cloud_storage.GetIfChanged(generated_profile_archive_path,
127          cloud_storage.PUBLIC_BUCKET)
128    except (cloud_storage.CredentialsError,
129            cloud_storage.PermissionError) as e:
130      if os.path.exists(generated_profile_archive_path):
131        # If the profile directory archive exists, assume the user has their
132        # own local copy simply warn.
133        logging.warning('Could not download Profile archive: %s',
134            generated_profile_archive_path)
135      else:
136        # If the archive profile directory doesn't exist, this is fatal.
137        logging.error('Can not run without required profile archive: %s. '
138                      'If you believe you have credentials, follow the '
139                      'instructions below.',
140                      generated_profile_archive_path)
141        logging.error(str(e))
142        sys.exit(-1)
143
144    # Unzip profile directory.
145    extracted_profile_dir_path = (
146        os.path.splitext(generated_profile_archive_path)[0])
147    if not os.path.isfile(generated_profile_archive_path):
148      raise Exception("Profile directory archive not downloaded: ",
149          generated_profile_archive_path)
150    with zipfile.ZipFile(generated_profile_archive_path) as f:
151      try:
152        f.extractall(os.path.dirname(generated_profile_archive_path))
153      except e:
154        # Cleanup any leftovers from unzipping.
155        if os.path.exists(extracted_profile_dir_path):
156          shutil.rmtree(extracted_profile_dir_path)
157        logging.error("Error extracting profile directory zip file: %s", e)
158        sys.exit(-1)
159
160    # Run with freshly extracted profile directory.
161    logging.info("Using profile archive directory: %s",
162        extracted_profile_dir_path)
163    options.browser_options.profile_dir = extracted_profile_dir_path
164
165  @classmethod
166  def PageTestClass(cls):
167    """Get the PageTest for this Benchmark.
168
169    If the Benchmark has no PageTest, raises NotImplementedError.
170    """
171    if not hasattr(cls, 'test'):
172      raise NotImplementedError('This test has no "test" attribute.')
173    if not issubclass(cls.test, page_test.PageTest):
174      raise TypeError('"%s" is not a PageTest.' % cls.test.__name__)
175    return cls.test
176
177  @classmethod
178  def PageSetClass(cls):
179    """Get the PageSet for this Benchmark.
180
181    If the Benchmark has no PageSet, raises NotImplementedError.
182    """
183    if not hasattr(cls, 'page_set'):
184      raise NotImplementedError('This test has no "page_set" attribute.')
185    if not issubclass(cls.page_set, page_set.PageSet):
186      raise TypeError('"%s" is not a PageSet.' % cls.page_set.__name__)
187    return cls.page_set
188
189  @classmethod
190  def CreatePageSet(cls, options):  # pylint: disable=W0613
191    """Get the page set this test will run on.
192
193    By default, it will create a page set from the file at this test's
194    page_set attribute. Override to generate a custom page set.
195    """
196    return cls.PageSetClass()()
197
198  @classmethod
199  def CreateExpectations(cls, ps):  # pylint: disable=W0613
200    """Get the expectations this test will run with.
201
202    By default, it will create an empty expectations set. Override to generate
203    custom expectations.
204    """
205    return test_expectations.TestExpectations()
206
207
208def AddCommandLineArgs(parser):
209  page_runner.AddCommandLineArgs(parser)
210
211
212def ProcessCommandLineArgs(parser, args):
213  page_runner.ProcessCommandLineArgs(parser, args)
214