update tests to use filetype instead of imghdr

Author: navin772

py/test/selenium/webdriver/common/takes_screenshots_tests.py

@@ -16,7 +16,7 @@
 # under the License.
 
 import base64
-import imghdr
+import filetype
 
 import pytest
 
@@ -26,13 +26,15 @@
 def test_get_screenshot_as_base64(driver, pages):
     pages.load("simpleTest.html")
     result = base64.b64decode(driver.get_screenshot_as_base64())
-    assert imghdr.what("", result) == "png"
+    kind = filetype.guess(result)
+    assert kind is not None and kind.mime == "image/png"
 
 
 def test_get_screenshot_as_png(driver, pages):
     pages.load("simpleTest.html")
     result = driver.get_screenshot_as_png()
-    assert imghdr.what("", result) == "png"
+    kind = filetype.guess(result)
+    assert kind is not None and kind.mime == "image/png"
 
 
 @pytest.mark.xfail_firefox
@@ -41,4 +43,5 @@ def test_get_element_screenshot(driver, pages):
     pages.load("simpleTest.html")
     element = driver.find_element(By.ID, "multiline")
     result = base64.b64decode(element.screenshot_as_base64)
-    assert imghdr.what("", result) == "png"
+    kind = filetype.guess(result)
+    assert kind is not None and kind.mime == "image/png"

py/test/selenium/webdriver/firefox/ff_takes_full_page_screenshots_tests.py

@@ -16,16 +16,18 @@
 # under the License.
 
 import base64
-import imghdr
+import filetype
 
 
 def test_get_full_page_screenshot_as_base64(driver, pages):
     pages.load("simpleTest.html")
     result = base64.b64decode(driver.get_full_page_screenshot_as_base64())
-    assert imghdr.what("", result) == "png"
+    kind = filetype.guess(result)
+    assert kind is not None and kind.mime == "image/png"
 
 
 def test_get_full_page_screenshot_as_png(driver, pages):
     pages.load("simpleTest.html")
     result = driver.get_full_page_screenshot_as_png()
-    assert imghdr.what("", result) == "png"
+    kind = filetype.guess(result)
+    assert kind is not None and kind.mime == "image/png"

py/test/selenium/webdriver/remote/remote_connection_tests.py

@@ -15,11 +15,12 @@
 # specific language governing permissions and limitations
 # under the License.
 import base64
-import imghdr
+import filetype
 
 
 def test_browser_specific_method(driver, pages):
     pages.load("simpleTest.html")
     screenshot = driver.execute("FULL_PAGE_SCREENSHOT")["value"]
     result = base64.b64decode(screenshot)
-    assert imghdr.what("", result) == "png"
+    kind = filetype.guess(result)
+    assert kind is not None and kind.mime == "image/png"