1# Copyright 2016 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 mock
6import tempfile
7import unittest
8import os
9
10from py_utils import cloud_storage  # pylint: disable=import-error
11from telemetry.testing import fakes
12from telemetry.internal.util import file_handle
13from telemetry.util import image_util
14from telemetry.util import screenshot
15
16class ScreenshotUtilTests(unittest.TestCase):
17
18  def setUp(self):
19    self.options = fakes.CreateBrowserFinderOptions()
20
21  def testScreenShotTakenSupportedPlatform(self):
22    fake_platform = self.options.fake_possible_browser.returned_browser.platform
23    expected_png_base64 = """
24      iVBORw0KGgoAAAANSUhEUgAAAAIAAAACCAIAAAD91
25      JpzAAAAFklEQVR4Xg3EAQ0AAABAMP1LY3YI7l8l6A
26      T8tgwbJAAAAABJRU5ErkJggg==
27      """
28    fake_platform.screenshot_png_data = expected_png_base64
29
30    fh = screenshot.TryCaptureScreenShot(fake_platform, None)
31    screenshot_file_path = fh.GetAbsPath()
32    try:
33      actual_screenshot_img = image_util.FromPngFile(screenshot_file_path)
34      self.assertTrue(image_util.AreEqual(
35                      image_util.FromBase64Png(expected_png_base64),
36                      actual_screenshot_img))
37    finally:  # Must clean up screenshot file if exists.
38      os.remove(screenshot_file_path)
39
40  def testUploadScreenshotToCloudStorage(self):
41    tf = tempfile.NamedTemporaryFile(
42        suffix='.png', delete=False)
43    fh1 = file_handle.FromTempFile(tf)
44
45    local_path = '123456abcdefg.png'
46
47    with mock.patch('py_utils.cloud_storage.Insert') as mock_insert:
48      with mock.patch('telemetry.util.screenshot._GenerateRemotePath',
49        return_value=local_path):
50
51        url = screenshot._UploadScreenShotToCloudStorage(fh1)
52        mock_insert.assert_called_with(
53            cloud_storage.TELEMETRY_OUTPUT,
54            local_path,
55            fh1.GetAbsPath())
56        self.assertTrue(url is not None)
57