tor-browser

The Tor Browser
git clone https://git.dasho.dev/tor-browser.git
Log | Files | Refs | README | LICENSE

test_screenshot.py (10684B)


      1 # This Source Code Form is subject to the terms of the Mozilla Public
      2 # License, v. 2.0. If a copy of the MPL was not distributed with this
      3 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
      4 
      5 import base64
      6 import hashlib
      7 import struct
      8 import tempfile
      9 import unittest
     10 
     11 from urllib.parse import quote
     12 
     13 import mozinfo
     14 
     15 from marionette_driver import By
     16 from marionette_driver.errors import NoSuchWindowException
     17 from marionette_harness import (
     18    MarionetteTestCase,
     19    skip,
     20    WindowManagerMixin,
     21 )
     22 
     23 
     24 def inline(doc, mime="text/html;charset=utf-8"):
     25    return "data:{0},{1}".format(mime, quote(doc))
     26 
     27 
     28 box = inline(
     29    "<body><div id='box'><p id='green' style='width: 50px; height: 50px; "
     30    "background: silver;'></p></div></body>"
     31 )
     32 input = inline("<body><input id='text-input'></input></body>")
     33 long = inline("<body style='height: 300vh'><p style='margin-top: 100vh'>foo</p></body>")
     34 short = inline("<body style='height: 10vh'></body>")
     35 svg = inline(
     36    """
     37    <svg xmlns="http://www.w3.org/2000/svg" height="20" width="20">
     38      <rect height="20" width="20"/>
     39    </svg>""",
     40    mime="image/svg+xml",
     41 )
     42 
     43 
     44 class ScreenCaptureTestCase(MarionetteTestCase):
     45    def setUp(self):
     46        super(ScreenCaptureTestCase, self).setUp()
     47 
     48        self.maxDiff = None
     49 
     50        self._device_pixel_ratio = None
     51 
     52        # Ensure that each screenshot test runs on a blank page to avoid left
     53        # over elements or focus which could interfer with taking screenshots
     54        self.marionette.navigate("about:blank")
     55 
     56    @property
     57    def device_pixel_ratio(self):
     58        if self._device_pixel_ratio is None:
     59            self._device_pixel_ratio = self.marionette.execute_script(
     60                """
     61                return window.devicePixelRatio
     62                """
     63            )
     64        return self._device_pixel_ratio
     65 
     66    @property
     67    def document_element(self):
     68        return self.marionette.find_element(By.CSS_SELECTOR, ":root")
     69 
     70    @property
     71    def page_y_offset(self):
     72        return self.marionette.execute_script("return window.pageYOffset")
     73 
     74    @property
     75    def viewport_dimensions(self):
     76        return self.marionette.execute_script(
     77            "return [window.innerWidth, window.innerHeight];"
     78        )
     79 
     80    def assert_png(self, screenshot):
     81        """Test that screenshot is a Base64 encoded PNG file."""
     82        if not isinstance(screenshot, bytes):
     83            screenshot = bytes(screenshot, encoding="utf-8")
     84            image = base64.decodebytes(screenshot)
     85        else:
     86            if screenshot.startswith(b"\211PNG\r\n\032\n"):
     87                image = screenshot
     88            else:
     89                image = base64.decodebytes(screenshot)
     90        self.assertRegex(image, b"\211PNG\r\n\032\n", "Expected image to be PNG")
     91        return image
     92 
     93    def assert_formats(self, element=None):
     94        if element is None:
     95            element = self.document_element
     96 
     97        image_default = self.assert_png(self.marionette.screenshot(element=element))
     98        screenshot_base64 = self.marionette.screenshot(element=element, format="base64")
     99        image_base64 = self.assert_png(screenshot_base64)
    100        image_binary1 = self.marionette.screenshot(element=element, format="binary")
    101        image_binary2 = self.marionette.screenshot(element=element, format="binary")
    102        screenshot_hash1 = self.marionette.screenshot(element=element, format="hash")
    103        screenshot_hash2 = self.marionette.screenshot(element=element, format="hash")
    104 
    105        # Valid data should have been returned
    106        self.assert_png(image_base64)
    107        self.assert_png(image_binary1)
    108        self.assertEqual(image_base64, image_binary1)
    109        self.assertEqual(
    110            screenshot_hash1,
    111            hashlib.sha256(screenshot_base64.encode("utf-8")).hexdigest(),
    112        )
    113 
    114        # Different formats produce different data
    115        self.assertNotEqual(screenshot_base64, image_binary1)
    116        self.assertNotEqual(screenshot_base64, screenshot_hash1)
    117        self.assertNotEqual(image_binary1, screenshot_hash1)
    118 
    119        # A second capture should be identical
    120        self.assertEqual(image_base64, image_default)
    121        self.assertEqual(image_binary1, image_binary2)
    122        self.assertEqual(screenshot_hash1, screenshot_hash2)
    123 
    124    def get_element_dimensions(self, element):
    125        rect = element.rect
    126        return rect["width"], rect["height"]
    127 
    128    def get_image_dimensions(self, image):
    129        image = self.assert_png(image)
    130        width, height = struct.unpack(">LL", image[16:24])
    131        return int(width), int(height)
    132 
    133    def scale(self, rect):
    134        return (
    135            int(rect[0] * self.device_pixel_ratio),
    136            int(rect[1] * self.device_pixel_ratio),
    137        )
    138 
    139 
    140 class TestScreenCaptureContent(WindowManagerMixin, ScreenCaptureTestCase):
    141    def setUp(self):
    142        super(TestScreenCaptureContent, self).setUp()
    143        self.marionette.set_context("content")
    144 
    145    def tearDown(self):
    146        self.close_all_tabs()
    147        super(TestScreenCaptureContent, self).tearDown()
    148 
    149    @property
    150    def scroll_dimensions(self):
    151        return tuple(
    152            self.marionette.execute_script(
    153                """
    154            return [
    155              document.documentElement.scrollWidth,
    156              document.documentElement.scrollHeight
    157            ];
    158            """
    159            )
    160        )
    161 
    162    def test_capture_tab_already_closed(self):
    163        new_tab = self.open_tab()
    164        self.marionette.switch_to_window(new_tab)
    165        self.marionette.close()
    166 
    167        self.assertRaises(NoSuchWindowException, self.marionette.screenshot)
    168        self.marionette.switch_to_window(self.start_tab)
    169 
    170    @unittest.skipIf(mozinfo.info["bits"] == 32, "Bug 1582973 - Risk for OOM on 32bit")
    171    def test_capture_vertical_bounds(self):
    172        self.marionette.navigate(inline("<body style='margin-top: 32768px'>foo"))
    173        screenshot = self.marionette.screenshot()
    174        self.assert_png(screenshot)
    175 
    176    @unittest.skipIf(mozinfo.info["bits"] == 32, "Bug 1582973 - Risk for OOM on 32bit")
    177    def test_capture_horizontal_bounds(self):
    178        self.marionette.navigate(inline("<body style='margin-left: 32768px'>foo"))
    179        screenshot = self.marionette.screenshot()
    180        self.assert_png(screenshot)
    181 
    182    @unittest.skipIf(mozinfo.info["bits"] == 32, "Bug 1582973 - Risk for OOM on 32bit")
    183    def test_capture_area_bounds(self):
    184        self.marionette.navigate(
    185            inline("<body style='margin-right: 21747px; margin-top: 21747px'>foo")
    186        )
    187        screenshot = self.marionette.screenshot()
    188        self.assert_png(screenshot)
    189 
    190    def test_capture_element(self):
    191        self.marionette.navigate(box)
    192        el = self.marionette.find_element(By.TAG_NAME, "div")
    193        screenshot = self.marionette.screenshot(element=el)
    194        self.assert_png(screenshot)
    195        self.assertEqual(
    196            self.scale(self.get_element_dimensions(el)),
    197            self.get_image_dimensions(screenshot),
    198        )
    199 
    200    @skip("Bug 1213875")
    201    def test_capture_element_scrolled_into_view(self):
    202        self.marionette.navigate(long)
    203        el = self.marionette.find_element(By.TAG_NAME, "p")
    204        screenshot = self.marionette.screenshot(element=el)
    205        self.assert_png(screenshot)
    206        self.assertEqual(
    207            self.scale(self.get_element_dimensions(el)),
    208            self.get_image_dimensions(screenshot),
    209        )
    210        self.assertGreater(self.page_y_offset, 0)
    211 
    212    def test_capture_full_html_document_element(self):
    213        self.marionette.navigate(long)
    214        screenshot = self.marionette.screenshot()
    215        self.assert_png(screenshot)
    216        self.assertEqual(
    217            self.scale(self.scroll_dimensions), self.get_image_dimensions(screenshot)
    218        )
    219 
    220    def test_capture_full_svg_document_element(self):
    221        self.marionette.navigate(svg)
    222        screenshot = self.marionette.screenshot()
    223        self.assert_png(screenshot)
    224        self.assertEqual(
    225            self.scale(self.scroll_dimensions), self.get_image_dimensions(screenshot)
    226        )
    227 
    228    def test_capture_viewport(self):
    229        url = self.marionette.absolute_url("clicks.html")
    230        self.marionette.navigate(short)
    231        self.marionette.navigate(url)
    232        screenshot = self.marionette.screenshot(full=False)
    233        self.assert_png(screenshot)
    234        self.assertEqual(
    235            self.scale(self.viewport_dimensions), self.get_image_dimensions(screenshot)
    236        )
    237 
    238    def test_capture_viewport_after_scroll(self):
    239        self.marionette.navigate(long)
    240        before = self.marionette.screenshot()
    241        el = self.marionette.find_element(By.TAG_NAME, "p")
    242        self.marionette.execute_script(
    243            "arguments[0].scrollIntoView()", script_args=[el]
    244        )
    245        after = self.marionette.screenshot(full=False)
    246        self.assertNotEqual(before, after)
    247        self.assertGreater(self.page_y_offset, 0)
    248 
    249    def test_formats(self):
    250        self.marionette.navigate(box)
    251 
    252        # Use a smaller region to speed up the test
    253        element = self.marionette.find_element(By.TAG_NAME, "div")
    254        self.assert_formats(element=element)
    255 
    256    def test_format_unknown(self):
    257        with self.assertRaises(ValueError):
    258            self.marionette.screenshot(format="cheese")
    259 
    260    def test_save_screenshot(self):
    261        expected = self.marionette.screenshot(format="binary")
    262        with tempfile.TemporaryFile("w+b") as fh:
    263            self.marionette.save_screenshot(fh)
    264            fh.flush()
    265            fh.seek(0)
    266            content = fh.read()
    267            self.assertEqual(expected, content)
    268 
    269    def test_scroll_default(self):
    270        self.marionette.navigate(long)
    271        before = self.page_y_offset
    272        el = self.marionette.find_element(By.TAG_NAME, "p")
    273        self.marionette.screenshot(element=el, format="hash")
    274        self.assertNotEqual(before, self.page_y_offset)
    275 
    276    def test_scroll(self):
    277        self.marionette.navigate(long)
    278        before = self.page_y_offset
    279        el = self.marionette.find_element(By.TAG_NAME, "p")
    280        self.marionette.screenshot(element=el, format="hash", scroll=True)
    281        self.assertNotEqual(before, self.page_y_offset)
    282 
    283    def test_scroll_off(self):
    284        self.marionette.navigate(long)
    285        el = self.marionette.find_element(By.TAG_NAME, "p")
    286        before = self.page_y_offset
    287        self.marionette.screenshot(element=el, format="hash", scroll=False)
    288        self.assertEqual(before, self.page_y_offset)
    289 
    290    def test_scroll_no_element(self):
    291        self.marionette.navigate(long)
    292        before = self.page_y_offset
    293        self.marionette.screenshot(format="hash", scroll=True)
    294        self.assertEqual(before, self.page_y_offset)