[I-Spy] Add support for rebaselining expectations from the web UI.

chrome/test/functional/ispy/common/chrome_utils.py

Index: chrome/test/functional/ispy/common/chrome_utils.py
diff --git a/chrome/test/functional/ispy/client/chrome_utils.py b/chrome/test/functional/ispy/common/chrome_utils.py
similarity index 53%
rename from chrome/test/functional/ispy/client/chrome_utils.py
rename to chrome/test/functional/ispy/common/chrome_utils.py
index 59166f650ea67ce21744c8b8ac48313b1fb4459d..10f1a9d832f6ccdb84a18c926063fcacb3fbd746 100644
--- a/chrome/test/functional/ispy/client/chrome_utils.py
+++ b/chrome/test/functional/ispy/common/chrome_utils.py
@@ -5,48 +5,42 @@
 import json
 import logging
 import os
-import time
 from distutils.version import LooseVersion
 from PIL import Image
 
-from ..common import cloud_bucket
-from ..common import ispy_utils
+import cloud_bucket
+import ispy_utils
 
 
 class ChromeUtils(object):
   """A utility for using ISpy with Chrome."""
 
-  def __init__(self, cloud_bucket, version_file, screenshot_func):
+  def __init__(self, cloud_bucket):
     """Initializes the utility class.
 
     Args:
       cloud_bucket: a BaseCloudBucket in which to the version file,
           expectations and results are to be stored.
-      version_file: path to the version file in the cloud bucket. The version
-          file contains a json list of ordered Chrome versions for which
-          expectations exist.
-      screenshot_func: a function that returns a PIL.Image.
     """
     self._cloud_bucket = cloud_bucket
-    self._version_file = version_file
-    self._screenshot_func = screenshot_func
     self._ispy = ispy_utils.ISpyUtils(self._cloud_bucket)
-    with open(
-        os.path.join(os.path.dirname(__file__), 'wait_on_ajax.js'), 'r') as f:
-      self._wait_for_unchanging_dom_script = f.read()
+    self._rebaselineable_cache = {}
 
-  def UpdateExpectationVersion(self, chrome_version):
+  def UpdateExpectationVersion(self, chrome_version, version_file):
     """Updates the most recent expectation version to the Chrome version.
 
     Should be called after generating a new set of expectations.
 
     Args:
       chrome_version: the chrome version as a string of the form "31.0.123.4".
+      version_file: path to the version file in the cloud bucket. The version
+          file contains a json list of ordered Chrome versions for which
+          expectations exist.
     """
     insert_pos = 0
     expectation_versions = []
     try:
-      expectation_versions = self._GetExpectationVersionList()
+      expectation_versions = self._GetExpectationVersionList(version_file)
       if expectation_versions:
         try:
           version = self._GetExpectationVersion(
@@ -61,7 +55,7 @@ class ChromeUtils(object):
     expectation_versions.insert(insert_pos, chrome_version)
     logging.info('Updating expectation version...')
     self._cloud_bucket.UploadFile(
-        self._version_file, json.dumps(expectation_versions),
+        version_file, json.dumps(expectation_versions),
         'application/json')
 
   def _GetExpectationVersion(self, chrome_version, expectation_versions):
@@ -81,12 +75,21 @@ class ChromeUtils(object):
         return version
     raise Exception('No expectation exists for Chrome %s' % chrome_version)
 
-  def _GetExpectationVersionList(self):
-    """Gets the list of expectation versions from google storage."""
-    return json.loads(self._cloud_bucket.DownloadFile(self._version_file))
+  def _GetExpectationVersionList(self, version_file):
+    """Gets the list of expectation versions from google storage.
+
+    Args:
+      version_file: path to the version file in the cloud bucket. The version
+          file contains a json list of ordered Chrome versions for which
+          expectations exist.
+
+    Returns:
+      Ordered list of Chrome versions.
+    """
+    return json.loads(self._cloud_bucket.DownloadFile(version_file))
 
   def _GetExpectationNameWithVersion(self, device_type, expectation,
-                                     chrome_version):
+                                     chrome_version, version_file):
     """Get the expectation to be used with the current Chrome version.
 
     Args:
@@ -98,7 +101,7 @@ class ChromeUtils(object):
       Version as an integer.
     """
     version = self._GetExpectationVersion(
-        chrome_version, self._GetExpectationVersionList())
+        chrome_version, self._GetExpectationVersionList(version_file))
     return self._CreateExpectationName(device_type, expectation, version)
 
   def _CreateExpectationName(self, device_type, expectation, version):
@@ -114,47 +117,110 @@ class ChromeUtils(object):
     """
     return '%s:%s(%s)' % (device_type, expectation, version)
 
-  def GenerateExpectation(self, device_type, expectation, chrome_version):
-    """Take screenshots and store as an expectation in I-Spy.
+  def GenerateExpectation(self, device_type, expectation, chrome_version,
+                          version_file, screenshots):
+    """Create an expectation for I-Spy.
 
     Args:
       device_type: string identifier for the device type.
       expectation: name for the expectation to generate.
       chrome_version: the chrome version as a string of the form "31.0.123.4".
+      screenshots: a list of similar PIL.Images.
     """
     # https://code.google.com/p/chromedriver/issues/detail?id=463
-    time.sleep(1)
     expectation_with_version = self._CreateExpectationName(
         device_type, expectation, chrome_version)
     if self._ispy.ExpectationExists(expectation_with_version):
       logging.warning(
           'I-Spy expectation \'%s\' already exists, overwriting.',
           expectation_with_version)
-    screenshots = [self._screenshot_func() for _ in range(8)]
     logging.info('Generating I-Spy expectation...')
     self._ispy.GenerateExpectation(expectation_with_version, screenshots)
 
   def PerformComparison(self, test_run, device_type, expectation,
-                        chrome_version):
-    """Take a screenshot and compare it with the given expectation in I-Spy.
+                        chrome_version, version_file, screenshot):
+    """Compare a screenshot with the given expectation in I-Spy.
 
     Args:
       test_run: name for the test run.
       device_type: string identifier for the device type.
       expectation: name for the expectation to compare against.
       chrome_version: the chrome version as a string of the form "31.0.123.4".
+      screenshot: a PIL.Image to compare.
     """
     # https://code.google.com/p/chromedriver/issues/detail?id=463
-    time.sleep(1)
-    screenshot = self._screenshot_func()
     logging.info('Performing I-Spy comparison...')
     self._ispy.PerformComparison(
         test_run,
         self._GetExpectationNameWithVersion(
-            device_type, expectation, chrome_version),
+            device_type, expectation, chrome_version, version_file),
         screenshot)
 
-  def GetScriptToWaitForUnchangingDOM(self):
-    """Returns a JavaScript script that waits for the DOM to stop changing."""
-    return self._wait_for_unchanging_dom_script
+  def CanRebaselineToTestRun(self, test_run):
+    """Returns whether the test run has associated expectations.
+
+    Returns:
+      True if RebaselineToTestRun() can be called for this test run.
+    """
+    if test_run in self._rebaselineable_cache:
+      return True
+    return self._cloud_bucket.FileExists(
+        ispy_utils.GetTestRunPath(test_run, 'rebaseline.txt'))
+
+  def RebaselineToTestRun(self, test_run):
+    """Update the version file to use expectations associated with |test_run|.
+
+    Args:
+      test_run: The name of the test run to rebaseline.
+    """
+    rebaseline_path = ispy_utils.GetTestRunPath(test_run, 'rebaseline.txt')
+    rebaseline_attrib = json.loads(
+        self._cloud_bucket.DownloadFile(rebaseline_path))
+    self.UpdateExpectationVersion(
+        rebaseline_attrib['version'], rebaseline_attrib['version_file'])
+    self._cloud_bucket.RemoveFile(rebaseline_path)
+
+  def _SetTestRunRebaselineable(self, test_run, chrome_version, version_file):
+    """Writes a JSON file containing the data needed to rebaseline.
+
+    Args:
+      test_run: The name of the test run to add the rebaseline file to.
+      chrome_version: the chrome version that can be rebaselined to (must have
+        associated Expectations).
+      version_file: the path of the version file associated with the test run.
+    """
+    self._rebaselineable_cache[test_run] = True
+    self._cloud_bucket.UploadFile(
+        ispy_utils.GetTestRunPath(test_run, 'rebaseline.txt'),
+        json.dumps({
+            'version': chrome_version,
+            'version_file': version_file}),
+        'application/json')
+
+  def PerformComparisonAndPrepareExpectation(self, test_run, device_type,
+                                             expectation, chrome_version,
+                                             version_file, screenshots):
+    """Perform comparison and generate an expectation that can used later.
+
+    The test run web UI will have a button to set the Expectations generated for
+    this version as the expectation for comparison with later versions.
+
+    Args:
+      test_run: The name of the test run to add the rebaseline file to.
+      device_type: string identifier for the device type.
+      chrome_version: the chrome version that can be rebaselined to (must have
+        associated Expectations).
+      version_file: the path of the version file associated with the test run.
+      screenshot: a list of similar PIL.Images.
+    """
+    if not self.CanRebaselineToTestRun(test_run):
+      self._SetTestRunRebaselineable(test_run, chrome_version, version_file)
+    expectation_with_version = self._CreateExpectationName(
+        device_type, expectation, chrome_version)
+    self._ispy.GenerateExpectation(expectation_with_version, screenshots)
+    self._ispy.PerformComparison(
+        test_run,
+        self._GetExpectationNameWithVersion(
+            device_type, expectation, chrome_version, version_file),
+        screenshots[-1])