# This Source Code Form is subject to the terms of the Mozilla Public # License, v. 2.0. If a copy of the MPL was not distributed with this # file, You can obtain one at http://mozilla.org/MPL/2.0/. import base64 import hashlib import struct import tempfile import unittest from urllib.parse import quote import mozinfo from marionette_driver import By from marionette_driver.errors import NoSuchWindowException from marionette_harness import ( MarionetteTestCase, skip, WindowManagerMixin, ) def inline(doc, mime="text/html;charset=utf-8"): return "data:{0},{1}".format(mime, quote(doc)) box = inline( "
foo
") short = inline("") svg = inline( """ """, mime="image/svg+xml", ) class ScreenCaptureTestCase(MarionetteTestCase): def setUp(self): super(ScreenCaptureTestCase, self).setUp() self.maxDiff = None self._device_pixel_ratio = None # Ensure that each screenshot test runs on a blank page to avoid left # over elements or focus which could interfer with taking screenshots self.marionette.navigate("about:blank") @property def device_pixel_ratio(self): if self._device_pixel_ratio is None: self._device_pixel_ratio = self.marionette.execute_script( """ return window.devicePixelRatio """ ) return self._device_pixel_ratio @property def document_element(self): return self.marionette.find_element(By.CSS_SELECTOR, ":root") @property def page_y_offset(self): return self.marionette.execute_script("return window.pageYOffset") @property def viewport_dimensions(self): return self.marionette.execute_script( "return [window.innerWidth, window.innerHeight];" ) def assert_png(self, screenshot): """Test that screenshot is a Base64 encoded PNG file.""" if not isinstance(screenshot, bytes): screenshot = bytes(screenshot, encoding="utf-8") image = base64.decodebytes(screenshot) else: if screenshot.startswith(b"\211PNG\r\n\032\n"): image = screenshot else: image = base64.decodebytes(screenshot) self.assertRegex(image, b"\211PNG\r\n\032\n", "Expected image to be PNG") return image def assert_formats(self, element=None): if element is None: element = self.document_element image_default = self.assert_png(self.marionette.screenshot(element=element)) screenshot_base64 = self.marionette.screenshot(element=element, format="base64") image_base64 = self.assert_png(screenshot_base64) image_binary1 = self.marionette.screenshot(element=element, format="binary") image_binary2 = self.marionette.screenshot(element=element, format="binary") screenshot_hash1 = self.marionette.screenshot(element=element, format="hash") screenshot_hash2 = self.marionette.screenshot(element=element, format="hash") # Valid data should have been returned self.assert_png(image_base64) self.assert_png(image_binary1) self.assertEqual(image_base64, image_binary1) self.assertEqual( screenshot_hash1, hashlib.sha256(screenshot_base64.encode("utf-8")).hexdigest(), ) # Different formats produce different data self.assertNotEqual(screenshot_base64, image_binary1) self.assertNotEqual(screenshot_base64, screenshot_hash1) self.assertNotEqual(image_binary1, screenshot_hash1) # A second capture should be identical self.assertEqual(image_base64, image_default) self.assertEqual(image_binary1, image_binary2) self.assertEqual(screenshot_hash1, screenshot_hash2) def get_element_dimensions(self, element): rect = element.rect return rect["width"], rect["height"] def get_image_dimensions(self, image): image = self.assert_png(image) width, height = struct.unpack(">LL", image[16:24]) return int(width), int(height) def scale(self, rect): return ( int(rect[0] * self.device_pixel_ratio), int(rect[1] * self.device_pixel_ratio), ) class TestScreenCaptureContent(WindowManagerMixin, ScreenCaptureTestCase): def setUp(self): super(TestScreenCaptureContent, self).setUp() self.marionette.set_context("content") def tearDown(self): self.close_all_tabs() super(TestScreenCaptureContent, self).tearDown() @property def scroll_dimensions(self): return tuple( self.marionette.execute_script( """ return [ document.documentElement.scrollWidth, document.documentElement.scrollHeight ]; """ ) ) def test_capture_tab_already_closed(self): new_tab = self.open_tab() self.marionette.switch_to_window(new_tab) self.marionette.close() self.assertRaises(NoSuchWindowException, self.marionette.screenshot) self.marionette.switch_to_window(self.start_tab) @unittest.skipIf(mozinfo.info["bits"] == 32, "Bug 1582973 - Risk for OOM on 32bit") def test_capture_vertical_bounds(self): self.marionette.navigate(inline("foo")) screenshot = self.marionette.screenshot() self.assert_png(screenshot) @unittest.skipIf(mozinfo.info["bits"] == 32, "Bug 1582973 - Risk for OOM on 32bit") def test_capture_horizontal_bounds(self): self.marionette.navigate(inline("foo")) screenshot = self.marionette.screenshot() self.assert_png(screenshot) @unittest.skipIf(mozinfo.info["bits"] == 32, "Bug 1582973 - Risk for OOM on 32bit") def test_capture_area_bounds(self): self.marionette.navigate( inline("foo") ) screenshot = self.marionette.screenshot() self.assert_png(screenshot) def test_capture_element(self): self.marionette.navigate(box) el = self.marionette.find_element(By.TAG_NAME, "div") screenshot = self.marionette.screenshot(element=el) self.assert_png(screenshot) self.assertEqual( self.scale(self.get_element_dimensions(el)), self.get_image_dimensions(screenshot), ) @skip("Bug 1213875") def test_capture_element_scrolled_into_view(self): self.marionette.navigate(long) el = self.marionette.find_element(By.TAG_NAME, "p") screenshot = self.marionette.screenshot(element=el) self.assert_png(screenshot) self.assertEqual( self.scale(self.get_element_dimensions(el)), self.get_image_dimensions(screenshot), ) self.assertGreater(self.page_y_offset, 0) def test_capture_full_html_document_element(self): self.marionette.navigate(long) screenshot = self.marionette.screenshot() self.assert_png(screenshot) self.assertEqual( self.scale(self.scroll_dimensions), self.get_image_dimensions(screenshot) ) def test_capture_full_svg_document_element(self): self.marionette.navigate(svg) screenshot = self.marionette.screenshot() self.assert_png(screenshot) self.assertEqual( self.scale(self.scroll_dimensions), self.get_image_dimensions(screenshot) ) def test_capture_viewport(self): url = self.marionette.absolute_url("clicks.html") self.marionette.navigate(short) self.marionette.navigate(url) screenshot = self.marionette.screenshot(full=False) self.assert_png(screenshot) self.assertEqual( self.scale(self.viewport_dimensions), self.get_image_dimensions(screenshot) ) def test_capture_viewport_after_scroll(self): self.marionette.navigate(long) before = self.marionette.screenshot() el = self.marionette.find_element(By.TAG_NAME, "p") self.marionette.execute_script( "arguments[0].scrollIntoView()", script_args=[el] ) after = self.marionette.screenshot(full=False) self.assertNotEqual(before, after) self.assertGreater(self.page_y_offset, 0) def test_formats(self): self.marionette.navigate(box) # Use a smaller region to speed up the test element = self.marionette.find_element(By.TAG_NAME, "div") self.assert_formats(element=element) def test_format_unknown(self): with self.assertRaises(ValueError): self.marionette.screenshot(format="cheese") def test_save_screenshot(self): expected = self.marionette.screenshot(format="binary") with tempfile.TemporaryFile("w+b") as fh: self.marionette.save_screenshot(fh) fh.flush() fh.seek(0) content = fh.read() self.assertEqual(expected, content) def test_scroll_default(self): self.marionette.navigate(long) before = self.page_y_offset el = self.marionette.find_element(By.TAG_NAME, "p") self.marionette.screenshot(element=el, format="hash") self.assertNotEqual(before, self.page_y_offset) def test_scroll(self): self.marionette.navigate(long) before = self.page_y_offset el = self.marionette.find_element(By.TAG_NAME, "p") self.marionette.screenshot(element=el, format="hash", scroll=True) self.assertNotEqual(before, self.page_y_offset) def test_scroll_off(self): self.marionette.navigate(long) el = self.marionette.find_element(By.TAG_NAME, "p") before = self.page_y_offset self.marionette.screenshot(element=el, format="hash", scroll=False) self.assertEqual(before, self.page_y_offset) def test_scroll_no_element(self): self.marionette.navigate(long) before = self.page_y_offset self.marionette.screenshot(format="hash", scroll=True) self.assertEqual(before, self.page_y_offset)