Python Tools for Pixel-by-Pixel Image Comparison

chrome/test/functional/ispy/ispy_core/Tools/bucket_manager.py

Index: chrome/test/functional/ispy/ispy_core/Tools/bucket_manager.py
diff --git a/chrome/test/functional/ispy/ispy_core/Tools/bucket_manager.py b/chrome/test/functional/ispy/ispy_core/Tools/bucket_manager.py
new file mode 100644
index 0000000000000000000000000000000000000000..638cbe50be6f79ed53c5b910cc1b3a4aed38122f
--- /dev/null
+++ b/chrome/test/functional/ispy/ispy_core/Tools/bucket_manager.py
@@ -0,0 +1,236 @@
+# Copyright (c) 2013 The Chromium Authos. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Utilities for accessing data on Google Cloud Storage."""
+
+import StringIO
+import sys
+sys.path.append('gsutil/third_party/boto')
+from PIL import Image
+
+import image_tools
+
+KEY = 'GOOGMLY23XF2KSSJ6WHF'
+SECRET = 'YletkArpp6KSucfnz/kwfyCUnokm/MzJX3KeiB4Y'
+BUCKET = 'i-spy'
+
+
+class NotFoundException(Exception):
+  pass
+
+
+class Manager(object):
+  """Responsible for handling bucket operations through a provided injector."""
+
+  def __init__(self, boto_injector):
+    """Creates a bucket manager from a boto_injector.
+
+    Args:
+      bucket_injector: an instance of a class implementing
+        Tools.BucketManagerTests.DependancyInjection.boto_injector.
+
+    Returns:
+      an instance of Tools.bucket_manager.Manager
+    """
+    self.b = boto_injector
+    self.bucket = self._GetConnection()
+
+  def _AuthenticateBucket(self, bucket, key, secret):
+    """Returns a connection to a given bucket.
+
+    Args:
+      bucket: The name of the bucket to be accessed.
+      key: An access key for a cloudstorage project.
+      secret: An access secret for a cloudstorage project.
+
+    Returns:
+      an instance of boto.gs.connection, which is an access point
+        through which users can manipulate items in the given bucket.
+    """
+    uri = self.b.storage_uri('', 'gs')
+    conn = self.b.connect(uri, key, secret)
+    return self.b.get_bucket(conn, bucket)
+
+  def _GetConnection(self):
+    """Returns a connection to the i-spy bucket.
+
+    Returns:
+      an instance of boto.gs.bucket that serves as an access point
+        to the i-spy bucket.
+    """
+    return self._AuthenticateBucket(BUCKET, KEY, SECRET)
+
+  def MakeKeyToFile(self, path):
+    """Returns a key to a file that may or may not be in the i-spy bucket.
+
+    Args:
+      path: the path to generate a key for.
+
+    Returns:
+      an instance of boto.gs.key which is an access point to a particular
+        resource in the i-spy bucket.
+    """
+    k = self.b.get_key(self.bucket)
+    k.set_path(path)
+    return k
+
+  def GetKeyToFile(self, path):
+    """Returns a key to a file that necessarily exists in the i-spy bucket.
+
+    Args:
+      path: the path to the file to be retrieved.
+
+    Returns:
+      an instance of boto.gs.key, which is an access point to a particular
+        resource in the i-spy bucket.
+
+    Raises:
+      NotFoundException: if the file being queried could not be found.
+    """
+    k = self.MakeKeyToFile(path)
+    if k.exists():
+      return k
+    else:
+      raise NotFoundException('Could not get file.')
+
+  def UploadImage(self, full_path, image):
+    """Uploads an image to a specified path under the i-spy bucket.
+
+    This function takes a path and image as input, then stores the
+      data from the image in the provided path in the i-spy bucket.
+
+    Args:
+      full_path: the full path, including file extension, to where the
+        image should be stored.
+      image: the RGB image to be stored in the i-spy bucket.
+
+    Returns:
+      None.
+    """
+    key = self.MakeKeyToFile(full_path)
+    key.set_metadata('Content-Type', 'image/png')
+    key.set_contents_from_string(
+        image_tools.SerializeImage(image).decode('base64')
+        )
+
+  def UploadTest(self, test_name, images):
+    """Uploads a set of images as a test.
+
+    This function uploads a collection of images which are
+      accepted as being visually correct. A mask is then generated
+      to account for the differences between the images. This mask
+      will be used to occlude animated portions of the images in the
+      test so they do not throw-off the pixel-by-pixel comparison.
+      The mask is stored under /Tests/[test_name]_mask.png, and
+      the color image is stored as /Tests/[test_name]_expected.png.
+
+    Args:
+      test_name: the name to give the newly created test
+      images:  a list of images that will be processed into a mask,
+        and expected image.
+
+    Returns:
+      None.
+    """
+    mask = image_tools.CreateMask(images)
+    path = 'Tests/%s_%s.png'
+    self.UploadImage(path % (test_name, 'mask'), mask)
+    self.UploadImage(path % (test_name, 'expected'), images[0])
+
+  def UploadResult(self, test_name, run_name, expected,
+                   diff, actual):
+    """Uploads an expected, a difference, and an actual image as a result.
+
+    This function uploads a three images, expected, diff, and actual
+      to the i-spy bucket under /Results/[test_name]_[expected/diff/actual].png.
+
+    Args:
+      test_name: the name of the test that this result corresponds to.
+      run_name: the name of the particular run of the test to upload.
+      expected: the correct RGB image provided by the test.
+      diff: the difference between the correct RGB image and the image
+        that the test has been run on.
+      actual: the RGB image that the test was run on.
+
+    Returns:
+      None.
+    """
+    path = 'Results/' + test_name + '/' + run_name + '/%s.png'
+    self.UploadImage(path % 'expected', expected)
+    self.UploadImage(path % 'diff', diff)
+    self.UploadImage(path % 'actual', actual)
+
+  def GetTest(self, test_name):
+    """Returns a dictionary of the expected and mask images of a given test.
+
+    Args:
+      test_name: the name of the test to get.
+
+    Returns:
+      a dictionary of the expected, and mask RGB images of the test.
+
+    Raises:
+      NotFoundException: if the test being queried could not be found.
+    """
+    mask_key = self.GetKeyToFile('Tests/%s_mask.png' % test_name)
+    expected_key = self.GetKeyToFile('Tests/%s_expected.png' % test_name)
+    return {
+        'expected': Image.open(StringIO.StringIO(
+            expected_key.get_contents_as_string())),
+        'mask': Image.open(StringIO.StringIO(
+            mask_key.get_contents_as_string()))
+        }
+
+  def GetResult(self, test_name, run_name):
+    """Returns the Result of a particular run of a test.
+
+    This function returns a dictionary of three keys: 'expected',
+      'diff', and 'actual' which are mapped to urls pointing to
+      their respective images for the given result.
+
+    Args:
+      test_name: the name of the test to get a result from.
+      run_name: the name of a particular run of the given test.
+
+    Returns:
+      a dictionary of strings mapped to url-strings for the 'expected',
+        'diff', and 'actual' components of a result.
+
+    Raises:
+      NotFoundException: if the result being queried could not be found.
+    """
+    path = 'Results/%s/%s/' % (test_name, run_name) + '%s.png'
+    return {
+        'expected': self.GetURL(path % 'expected'),
+        'diff': self.GetURL(path % 'diff'),
+        'actual': self.GetURL(path % 'actual')
+        }
+
+  def GetURL(self, full_path):
+    """Returns the url to a file in the i-spy bucket.
+
+    Args:
+      full_path: the full path to the file.
+
+    Returns:
+      the url of the file as a string that will be valid for 5 minutes after
+        the function call.
+
+    Raises:
+      NotFoundException: if full_path does not exist in the ispy bucket.
+    """
+    key = self.GetKeyToFile(full_path)
+    return key.generate_url(300).replace('AWSAccessKeyId', 'GoogleAccessId')
+
+  def GetAllKeys(self, prefix=None):
+    """Returns all keys in the i-spy bucket starting with a prefix.
+
+    Args:
+      prefix: An optional prefix to get all keys starting with. If the prefix
+        is left empty, all keys in the bucket will be returned.
+
+    Returns:
+      a list of all the keys in the i-spy bucket.
+    """
+    return self.b.get_all_keys(self.bucket, prefix)