[wdspec] Use shared assert_png() method in screenshot tests.

Depends on D10626

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

bugzilla-url: https://bugzilla.mozilla.org/show_bug.cgi?id=1503804
gecko-commit: 8ee95012eaf0271baf406be0f0b712198f9e01da
gecko-integration-branch: autoland
gecko-reviewers: ato

webdriver/tests/take_element_screenshot/screenshot.py

@@ -1,7 +1,4 @@
-import base64
-import imghdr
-
-from tests.support.asserts import assert_error, assert_success
+from tests.support.asserts import assert_error, assert_png, assert_success
 from tests.support.inline import inline
 
 
@@ -27,8 +24,7 @@ def test_screenshot(session):
     response = take_element_screenshot(session, element.id)
     value = assert_success(response)
 
-    image = base64.decodestring(value)
-    assert imghdr.what("", image) == "png"
+    assert_png(value)
 
 
 def test_stale(session):

webdriver/tests/take_element_screenshot/user_prompts.py

@@ -1,11 +1,8 @@
 # META: timeout=long
 
-import base64
-import imghdr
-
 import pytest
 
-from tests.support.asserts import assert_success
+from tests.support.asserts import assert_png, assert_success
 from tests.support.inline import inline
 
 
@@ -30,8 +27,7 @@ def check_user_prompt_not_closed_without_exception(dialog_type):
         response = take_element_screenshot(session, element.id)
         value = assert_success(response)
 
-        image = base64.decodestring(value)
-        assert imghdr.what("", image) == "png"
+        assert_png(value)
 
         assert session.alert.text == dialog_type
         session.alert.dismiss()

webdriver/tests/take_screenshot/screenshot.py

@@ -1,7 +1,4 @@
-import base64
-import imghdr
-
-from tests.support.asserts import assert_error, assert_success
+from tests.support.asserts import assert_error, assert_png, assert_success
 from tests.support.inline import inline
 
 
@@ -21,5 +18,4 @@ def test_screenshot(session):
     response = take_screenshot(session)
     value = assert_success(response)
 
-    image = base64.decodestring(value)
-    assert imghdr.what("", image) == "png"
+    assert_png(value)

webdriver/tests/take_screenshot/user_prompts.py

@@ -1,11 +1,8 @@
 # META: timeout=long
 
-import base64
-import imghdr
-
 import pytest
 
-from tests.support.asserts import assert_success
+from tests.support.asserts import assert_png, assert_success
 from tests.support.inline import inline
 
 
@@ -24,8 +21,7 @@ def check_user_prompt_not_closed_without_exception(dialog_type):
         response = take_screenshot(session)
         value = assert_success(response)
 
-        image = base64.decodestring(value)
-        assert imghdr.what("", image) == "png"
+        assert_png(value)
 
         assert session.alert.text == dialog_type
         session.alert.dismiss()