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