tor-browser

commit 998d38948a9f6f1357e7b5ed5250c11358fa16b0
parent c932ecc4d8c782d19213cb23bbc1a9c4754a818d
Author: elkal98 <khalid.alhaddad98@gmail.com>
Date:   Thu,  6 Nov 2025 20:00:54 +0000

Bug 1991099 - [marionette] Remove deprecated imghdr module for screenshot unit tests. r=whimboo

Differential Revision: https://phabricator.services.mozilla.com/D270289

Diffstat:
M
testing/marionette/harness/marionette_harness/tests/unit/test_screenshot.py
 | 
58
++++++++++++++++++++++++++++++----------------------------

1 file changed, 30 insertions(+), 28 deletions(-)
diff --git a/testing/marionette/harness/marionette_harness/tests/unit/test_screenshot.py b/testing/marionette/harness/marionette_harness/tests/unit/test_screenshot.py
@@ -4,7 +4,6 @@
 
 import base64
 import hashlib
-import imghdr
 import struct
 import tempfile
 import unittest
@@ -82,49 +81,52 @@ class ScreenCaptureTestCase(MarionetteTestCase):
         """Test that screenshot is a Base64 encoded PNG file."""
         if not isinstance(screenshot, bytes):
             screenshot = bytes(screenshot, encoding="utf-8")
-        image = base64.decodebytes(screenshot)
-        self.assertEqual(imghdr.what("", image), "png")
+            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
 
-        screenshot_default = self.marionette.screenshot(element=element)
-        if not isinstance(screenshot_default, bytes):
-            screenshot_default = bytes(screenshot_default, encoding="utf-8")
-        screenshot_image = self.marionette.screenshot(element=element, format="base64")
-        if not isinstance(screenshot_image, bytes):
-            screenshot_image = bytes(screenshot_image, encoding="utf-8")
-        binary1 = self.marionette.screenshot(element=element, format="binary")
-        binary2 = self.marionette.screenshot(element=element, format="binary")
-        hash1 = self.marionette.screenshot(element=element, format="hash")
-        hash2 = self.marionette.screenshot(element=element, format="hash")
+        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(screenshot_image)
-        self.assertEqual(imghdr.what("", binary1), "png")
-        self.assertEqual(screenshot_image, base64.b64encode(binary1))
-        self.assertEqual(hash1, hashlib.sha256(screenshot_image).hexdigest())
+        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_image, binary1)
-        self.assertNotEqual(screenshot_image, hash1)
-        self.assertNotEqual(binary1, hash1)
+        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(screenshot_image, screenshot_default)
-        self.assertEqual(binary1, binary2)
-        self.assertEqual(hash1, hash2)
+        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, screenshot):
-        if not isinstance(screenshot, bytes):
-            screenshot = bytes(screenshot, encoding="utf-8")
-        self.assert_png(screenshot)
-        image = base64.decodebytes(screenshot)
+    def get_image_dimensions(self, image):
+        image = self.assert_png(image)
         width, height = struct.unpack(">LL", image[16:24])
         return int(width), int(height)