take screenshot again when first failed

atx/drivers/dummy.py

@@ -22,15 +22,16 @@ def __init__(self, *args, **kwargs):
         self._rotation = 1
         self.last_click = None
         self.serial = '1234'
+        self._fail_first_screenshot = False
 
-    @hook_wrap(consts.EVENT_SCREENSHOT)
-    def screenshot(self, filename=None):
+    def _take_screenshot(self):
         """ Take a screenshot """
         # screen size: 1280x720
+        if self._fail_first_screenshot:
+            self._fail_first_screenshot = False
+            raise IOError("dummy fail screenshot")
         screen_path = os.path.join(__dir__, '../../tests/media/dummy_screen.png')
         screen = Image.open(screen_path)
-        if filename:
-            screen.save(filename)
         return screen
 
     @property

atx/drivers/mixin.py

@@ -359,7 +359,7 @@ def region_screenshot(self, filename=None):
         """Deprecated
         Take part of the screenshot
         """
-        warnings.warn("deprecated, use screenshot().crop(bounds) instead", DeprecationWarning)
+        # warnings.warn("deprecated, use screenshot().crop(bounds) instead", DeprecationWarning)
         screen = self.__last_screen if self.__keep_screen else self.screenshot()
         if self.bounds:
             screen = screen.crop(self.bounds)
@@ -383,13 +383,12 @@ def screenshot(self, filename=None):
         """
         if self.__keep_screen:
             return self.__last_screen
-        screen = self.__last_screen = self._take_screenshot()
-        # FIXME(ssx): just not take screenshot again
-        # try:
-        #     screen = self._take_screenshot()
-        # except IOError:
-        #     # try taks screenshot again
-        #     screen = self._take_screenshot()
+        try:
+            screen = self._take_screenshot()
+        except IOError:
+            # try taks screenshot again
+            log.warn("warning, screenshot failed [2/1], retry again")
+            screen = self._take_screenshot()
         self.__last_screen = screen
         if filename:
             save_dir = os.path.dirname(filename) or '.'

tests/test_dummy.py

@@ -3,17 +3,18 @@
 
 import os
 import time
+import mock
 
 import pytest
 import atx
-
 from PIL import Image
 
 d = atx.connect(platform='dummy')
 
 def setup_function(f):
     d.resolution = (1280, 720)
 
+
 def test_setget_resolution():
     assert d.resolution == (720, 1280)
 
@@ -28,14 +29,35 @@ def test_setget_resolution():
     with pytest.raises(TypeError):
         d.resolution = 720
     assert d.resolution == (200, 400)
-
+
+
 def teardown_function(f):
     print('teardown')
-
-def test_screenshot():
+
+
+def test_screenshot_normal():
     screen = d.screenshot()
     assert screen is not None
 
+
+def test_screenshot_first_fail():
+    d._fail_first_screenshot = True
+    screen = d.screenshot()
+    assert screen is not None
+
+
+def test_screenshot_always_fail():
+    import types
+    d = atx.connect(platform='dummy')
+
+    def raise_ioerror(self):
+        raise IOError('error of io')
+
+    d._take_screenshot = types.MethodType(raise_ioerror, d)
+    with pytest.raises(IOError):
+        d.screenshot()
+
+
 def test_hook_screenshot():
     called = [False]