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 os
6import tempfile
7import unittest
8
9from telemetry.core import util
10from telemetry.util import image_util
11from telemetry.util import rgba_color
12
13# This is a simple base64 encoded 2x2 PNG which contains, in order, a single
14# Red, Yellow, Blue, and Green pixel.
15test_png = """
16 iVBORw0KGgoAAAANSUhEUgAAAAIAAAACCAIAAAD91
17 JpzAAAAFklEQVR4Xg3EAQ0AAABAMP1LY3YI7l8l6A
18 T8tgwbJAAAAABJRU5ErkJggg==
19"""
20test_png_path = os.path.join(util.GetUnittestDataDir(), 'test_png.png')
21test_png_2_path = os.path.join(util.GetUnittestDataDir(), 'test_png_2.png')
22
23class ImageUtilTest(unittest.TestCase):
24  def testReadFromBase64Png(self):
25    bmp = image_util.FromBase64Png(test_png)
26
27    self.assertEquals(2, image_util.Width(bmp))
28    self.assertEquals(2, image_util.Height(bmp))
29
30    image_util.GetPixelColor(bmp, 0, 0).AssertIsRGB(255, 0, 0)
31    image_util.GetPixelColor(bmp, 1, 1).AssertIsRGB(0, 255, 0)
32    image_util.GetPixelColor(bmp, 0, 1).AssertIsRGB(0, 0, 255)
33    image_util.GetPixelColor(bmp, 1, 0).AssertIsRGB(255, 255, 0)
34
35  def testReadFromPngFile(self):
36    file_bmp = image_util.FromPngFile(test_png_path)
37
38    self.assertEquals(2, image_util.Width(file_bmp))
39    self.assertEquals(2, image_util.Height(file_bmp))
40
41    image_util.GetPixelColor(file_bmp, 0, 0).AssertIsRGB(255, 0, 0)
42    image_util.GetPixelColor(file_bmp, 1, 1).AssertIsRGB(0, 255, 0)
43    image_util.GetPixelColor(file_bmp, 0, 1).AssertIsRGB(0, 0, 255)
44    image_util.GetPixelColor(file_bmp, 1, 0).AssertIsRGB(255, 255, 0)
45
46  def testWritePngToPngFile(self):
47    orig = image_util.FromPngFile(test_png_path)
48    temp_file = tempfile.NamedTemporaryFile(suffix='.png').name
49    image_util.WritePngFile(orig, temp_file)
50    new_file = image_util.FromPngFile(temp_file)
51    self.assertTrue(image_util.AreEqual(orig, new_file, likely_equal=True))
52
53  def testWritePngWithoutPngSuffixThrows(self):
54    orig = image_util.FromPngFile(test_png_path)
55    temp_file = tempfile.NamedTemporaryFile().name
56    self.assertRaises(AssertionError, image_util.WritePngFile,
57                      orig, temp_file)
58
59  def testWriteCroppedBmpToPngFile(self):
60    pixels = [255, 0, 0, 255, 255, 0, 0, 0, 0,
61              255, 255, 0, 0, 255, 0, 0, 0, 0]
62    orig = image_util.FromRGBPixels(3, 2, pixels)
63    orig = image_util.Crop(orig, 0, 0, 2, 2)
64    temp_file = tempfile.NamedTemporaryFile(suffix='.png').name
65    image_util.WritePngFile(orig, temp_file)
66    new_file = image_util.FromPngFile(temp_file)
67    self.assertTrue(image_util.AreEqual(orig, new_file, likely_equal=True))
68
69  def testIsEqual(self):
70    bmp = image_util.FromBase64Png(test_png)
71    file_bmp = image_util.FromPngFile(test_png_path)
72    self.assertTrue(image_util.AreEqual(bmp, file_bmp, likely_equal=True))
73
74  def testDiff(self):
75    file_bmp = image_util.FromPngFile(test_png_path)
76    file_bmp_2 = image_util.FromPngFile(test_png_2_path)
77
78    diff_bmp = image_util.Diff(file_bmp, file_bmp)
79
80    self.assertEquals(2, image_util.Width(diff_bmp))
81    self.assertEquals(2, image_util.Height(diff_bmp))
82
83    image_util.GetPixelColor(diff_bmp, 0, 0).AssertIsRGB(0, 0, 0)
84    image_util.GetPixelColor(diff_bmp, 1, 1).AssertIsRGB(0, 0, 0)
85    image_util.GetPixelColor(diff_bmp, 0, 1).AssertIsRGB(0, 0, 0)
86    image_util.GetPixelColor(diff_bmp, 1, 0).AssertIsRGB(0, 0, 0)
87
88    diff_bmp = image_util.Diff(file_bmp, file_bmp_2)
89
90    self.assertEquals(3, image_util.Width(diff_bmp))
91    self.assertEquals(3, image_util.Height(diff_bmp))
92
93    image_util.GetPixelColor(diff_bmp, 0, 0).AssertIsRGB(0, 255, 255)
94    image_util.GetPixelColor(diff_bmp, 1, 1).AssertIsRGB(255, 0, 255)
95    image_util.GetPixelColor(diff_bmp, 0, 1).AssertIsRGB(255, 255, 0)
96    image_util.GetPixelColor(diff_bmp, 1, 0).AssertIsRGB(0, 0, 255)
97
98    image_util.GetPixelColor(diff_bmp, 0, 2).AssertIsRGB(255, 255, 255)
99    image_util.GetPixelColor(diff_bmp, 1, 2).AssertIsRGB(255, 255, 255)
100    image_util.GetPixelColor(diff_bmp, 2, 0).AssertIsRGB(255, 255, 255)
101    image_util.GetPixelColor(diff_bmp, 2, 1).AssertIsRGB(255, 255, 255)
102    image_util.GetPixelColor(diff_bmp, 2, 2).AssertIsRGB(255, 255, 255)
103
104  def testGetBoundingBox(self):
105    pixels = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
106              0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0,
107              0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
108    bmp = image_util.FromRGBPixels(4, 3, pixels)
109    box, count = image_util.GetBoundingBox(bmp, rgba_color.RgbaColor(1, 0, 0))
110    self.assertEquals(box, (1, 1, 2, 1))
111    self.assertEquals(count, 2)
112
113    box, count = image_util.GetBoundingBox(bmp, rgba_color.RgbaColor(0, 1, 0))
114    self.assertEquals(box, None)
115    self.assertEquals(count, 0)
116
117  def testCrop(self):
118    pixels = [0, 0, 0, 1, 0, 0, 2, 0, 0, 3, 0, 0,
119              0, 1, 0, 1, 1, 0, 2, 1, 0, 3, 1, 0,
120              0, 2, 0, 1, 2, 0, 2, 2, 0, 3, 2, 0]
121    bmp = image_util.FromRGBPixels(4, 3, pixels)
122    bmp = image_util.Crop(bmp, 1, 2, 2, 1)
123
124    self.assertEquals(2, image_util.Width(bmp))
125    self.assertEquals(1, image_util.Height(bmp))
126    image_util.GetPixelColor(bmp, 0, 0).AssertIsRGB(1, 2, 0)
127    image_util.GetPixelColor(bmp, 1, 0).AssertIsRGB(2, 2, 0)
128    self.assertEquals(image_util.Pixels(bmp), bytearray([1, 2, 0, 2, 2, 0]))
129