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 json
6import logging
7import os
8from distutils.version import LooseVersion
9from PIL import Image
10
11from common import cloud_bucket
12from common import ispy_utils
13
14
15class ISpyApi(object):
16  """The public API for interacting with ISpy."""
17
18  def __init__(self, cloud_bucket):
19    """Initializes the utility class.
20
21    Args:
22      cloud_bucket: a BaseCloudBucket in which to the version file,
23          expectations and results are to be stored.
24    """
25    self._cloud_bucket = cloud_bucket
26    self._ispy = ispy_utils.ISpyUtils(self._cloud_bucket)
27    self._rebaselineable_cache = {}
28
29  def UpdateExpectationVersion(self, chrome_version, version_file):
30    """Updates the most recent expectation version to the Chrome version.
31
32    Should be called after generating a new set of expectations.
33
34    Args:
35      chrome_version: the chrome version as a string of the form "31.0.123.4".
36      version_file: path to the version file in the cloud bucket. The version
37          file contains a json list of ordered Chrome versions for which
38          expectations exist.
39    """
40    insert_pos = 0
41    expectation_versions = []
42    try:
43      expectation_versions = self._GetExpectationVersionList(version_file)
44      if expectation_versions:
45        try:
46          version = self._GetExpectationVersion(
47              chrome_version, expectation_versions)
48          if version == chrome_version:
49            return
50          insert_pos = expectation_versions.index(version)
51        except:
52          insert_pos = len(expectation_versions)
53    except cloud_bucket.FileNotFoundError:
54      pass
55    expectation_versions.insert(insert_pos, chrome_version)
56    logging.info('Updating expectation version...')
57    self._cloud_bucket.UploadFile(
58        version_file, json.dumps(expectation_versions),
59        'application/json')
60
61  def _GetExpectationVersion(self, chrome_version, expectation_versions):
62    """Returns the expectation version for the given Chrome version.
63
64    Args:
65      chrome_version: the chrome version as a string of the form "31.0.123.4".
66      expectation_versions: Ordered list of Chrome versions for which
67        expectations exist, as stored in the version file.
68
69    Returns:
70      Expectation version string.
71    """
72    # Find the closest version that is not greater than the chrome version.
73    for version in expectation_versions:
74      if LooseVersion(version) <= LooseVersion(chrome_version):
75        return version
76    raise Exception('No expectation exists for Chrome %s' % chrome_version)
77
78  def _GetExpectationVersionList(self, version_file):
79    """Gets the list of expectation versions from google storage.
80
81    Args:
82      version_file: path to the version file in the cloud bucket. The version
83          file contains a json list of ordered Chrome versions for which
84          expectations exist.
85
86    Returns:
87      Ordered list of Chrome versions.
88    """
89    try:
90      return json.loads(self._cloud_bucket.DownloadFile(version_file))
91    except:
92      return []
93
94  def _GetExpectationNameWithVersion(self, device_type, expectation,
95                                     chrome_version, version_file):
96    """Get the expectation to be used with the current Chrome version.
97
98    Args:
99      device_type: string identifier for the device type.
100      expectation: name for the expectation to generate.
101      chrome_version: the chrome version as a string of the form "31.0.123.4".
102
103    Returns:
104      Version as an integer.
105    """
106    version = self._GetExpectationVersion(
107        chrome_version, self._GetExpectationVersionList(version_file))
108    return self._CreateExpectationName(device_type, expectation, version)
109
110  def _CreateExpectationName(self, device_type, expectation, version):
111    """Create the full expectation name from the expectation and version.
112
113    Args:
114      device_type: string identifier for the device type, example: mako
115      expectation: base name for the expectation, example: google.com
116      version: expectation version, example: 31.0.23.1
117
118    Returns:
119      Full expectation name as a string, example: mako:google.com(31.0.23.1)
120    """
121    return '%s:%s(%s)' % (device_type, expectation, version)
122
123  def GenerateExpectation(self, device_type, expectation, chrome_version,
124                          version_file, screenshots):
125    """Create an expectation for I-Spy.
126
127    Args:
128      device_type: string identifier for the device type.
129      expectation: name for the expectation to generate.
130      chrome_version: the chrome version as a string of the form "31.0.123.4".
131      screenshots: a list of similar PIL.Images.
132    """
133    # https://code.google.com/p/chromedriver/issues/detail?id=463
134    expectation_with_version = self._CreateExpectationName(
135        device_type, expectation, chrome_version)
136    if self._ispy.ExpectationExists(expectation_with_version):
137      logging.warning(
138          'I-Spy expectation \'%s\' already exists, overwriting.',
139          expectation_with_version)
140    logging.info('Generating I-Spy expectation...')
141    self._ispy.GenerateExpectation(expectation_with_version, screenshots)
142
143  def PerformComparison(self, test_run, device_type, expectation,
144                        chrome_version, version_file, screenshot):
145    """Compare a screenshot with the given expectation in I-Spy.
146
147    Args:
148      test_run: name for the test run.
149      device_type: string identifier for the device type.
150      expectation: name for the expectation to compare against.
151      chrome_version: the chrome version as a string of the form "31.0.123.4".
152      screenshot: a PIL.Image to compare.
153    """
154    # https://code.google.com/p/chromedriver/issues/detail?id=463
155    logging.info('Performing I-Spy comparison...')
156    self._ispy.PerformComparison(
157        test_run,
158        self._GetExpectationNameWithVersion(
159            device_type, expectation, chrome_version, version_file),
160        screenshot)
161
162  def CanRebaselineToTestRun(self, test_run):
163    """Returns whether the test run has associated expectations.
164
165    Returns:
166      True if RebaselineToTestRun() can be called for this test run.
167    """
168    if test_run in self._rebaselineable_cache:
169      return True
170    return self._cloud_bucket.FileExists(
171        ispy_utils.GetTestRunPath(test_run, 'rebaseline.txt'))
172
173  def RebaselineToTestRun(self, test_run):
174    """Update the version file to use expectations associated with |test_run|.
175
176    Args:
177      test_run: The name of the test run to rebaseline.
178    """
179    rebaseline_path = ispy_utils.GetTestRunPath(test_run, 'rebaseline.txt')
180    rebaseline_attrib = json.loads(
181        self._cloud_bucket.DownloadFile(rebaseline_path))
182    self.UpdateExpectationVersion(
183        rebaseline_attrib['version'], rebaseline_attrib['version_file'])
184    self._cloud_bucket.RemoveFile(rebaseline_path)
185
186  def _SetTestRunRebaselineable(self, test_run, chrome_version, version_file):
187    """Writes a JSON file containing the data needed to rebaseline.
188
189    Args:
190      test_run: The name of the test run to add the rebaseline file to.
191      chrome_version: the chrome version that can be rebaselined to (must have
192        associated Expectations).
193      version_file: the path of the version file associated with the test run.
194    """
195    self._rebaselineable_cache[test_run] = True
196    self._cloud_bucket.UploadFile(
197        ispy_utils.GetTestRunPath(test_run, 'rebaseline.txt'),
198        json.dumps({
199            'version': chrome_version,
200            'version_file': version_file}),
201        'application/json')
202
203  def PerformComparisonAndPrepareExpectation(self, test_run, device_type,
204                                             expectation, chrome_version,
205                                             version_file, screenshots):
206    """Perform comparison and generate an expectation that can used later.
207
208    The test run web UI will have a button to set the Expectations generated for
209    this version as the expectation for comparison with later versions.
210
211    Args:
212      test_run: The name of the test run to add the rebaseline file to.
213      device_type: string identifier for the device type.
214      chrome_version: the chrome version that can be rebaselined to (must have
215        associated Expectations).
216      version_file: the path of the version file associated with the test run.
217      screenshot: a list of similar PIL.Images.
218    """
219    if not self.CanRebaselineToTestRun(test_run):
220      self._SetTestRunRebaselineable(test_run, chrome_version, version_file)
221    expectation_with_version = self._CreateExpectationName(
222        device_type, expectation, chrome_version)
223    self._ispy.GenerateExpectation(expectation_with_version, screenshots)
224    self._ispy.PerformComparison(
225        test_run,
226        self._GetExpectationNameWithVersion(
227            device_type, expectation, chrome_version, version_file),
228        screenshots[-1])
229
230