Python Selenium - Capturing ScreenShot on Error

As part of the test framework, at the point of error, we'd like to capture as much information as possible, including taking screenshots. In order to do this using a Python Unittest / Selenium framework we did 2 things.  1) Create a screenshot utility function that can work both across a local or remote webdriver. 2) Create a base test where we modified the run method to take a screenshot prior to saving the error message, and

Creating our screen capture utility function

Selenium webdriver offers a couple methods for capturing screenshot.  One works well for local instances of webdriver, and the other while much slower is better to use for RemoteWebDriver as binary data can be unpredictable across scripting over the wire.

class ScreenShotUtil:
    "Screenshot Utility Class"

    @staticmethod
    def take_screenshot(webdriver, file_name="error.png"):
        """
        @param webdriver: WebDriver.
        @type webdriver: WebDriver
        @param file_name: Name to label this screenshot.
        @type file_name: str 
        """
        if isinstance(webdriver, remote.webdriver.WebDriver):
            # Get Screenshot over the wire as base64
            base64_data = webdriver.get_screenshot_as_base64()
            screenshot_data = base64.decodestring(base64_data)
            screenshot_file = open(filename, "w")
            screenshot_file.write(screenshot_data)
            screenshot_file.close()
        else:
            webdriver.save_screenshot(filename)

Adding our screenshot capture rule

In Java's JUnit, you can easily just create a MethodRule (or TestWatcher) which you can use to annotate your tests.  Because Python's UnitTest does not have such a feature to register call backs or event listeners to failed test results, short of implementing or using a 3rd party Test Runner, an easy way to do this is to just create a BaseTest where we override the 'run' method and insert these calls.  To do this,

1. Open the source code for unittest.TestCase.  You can find it here, http://sourceforge.net/projects/pyunit/
2. Create a BaseTest class that extends TestCase.
3. Override the run() method definition with the copy from the latest version.
4. Inside the exception hander for failed step, and error step, insert your call to your screen capture utility.

...
class ScreenCaptureTestCase(unittest.TestCase):
...
    # Defining an init method so we can pass it a webdriver.
    def __init__(self, methodName='runTest', webdriver=None, screenshot_util=None):
        super(WDBaseTest, self).__init__(methodName)
        
        if webdriver_provider == None:
            self._webdriver = WebDriverSingleton.get_instance()
        else:
            self._webdriver = webdriver

        if screenshot_util == None:
            self._screenshot_util = WebScreenShotUtil
        else:
            self._screenshot_util = screenshot_util 
    ...
def run(self, result=None):
        """
        Overriding the run() method to insert our screenshot handler.
        
        Most of this method is a copy of the TestCase.run() method source.
        """
        orig_result = result
        if result is None:
            result = self.defaultTestResult()
            startTestRun = getattr(result, 'startTestRun', None)
            if startTestRun is not None:
                startTestRun()
            ... more pyunit code ...
                except self.failureException:
                    # Insert our Take Screenshot on test failure code.
                    fname = str(self).replace("(", "").replace(")", "").replace(" ", "_")
                    fmt='%y-%m-%d_%H.%M.%S_.PNG'
                    filename = datetime.datetime.now().strftime(fmt)
                    self._screenshot_util.take_screenshot(self._webdriver, filename)
                    result.addFailure(self, sys.exc_info())
            ... more pyunit code...
                except:
                    # Do the same thing again for errors.
                    fname = str(self).replace("(", "").replace(")", "").replace(" ", "_")
                    fmt='%y-%m-%d_%H.%M.%S_.PNG'
                    filename = datetime.datetime.now().strftime(fmt)
                    self._screenshot_util.take_screenshot(self._webdriver, filename)
                    result.addError(self, sys.exc_info())
            ... more pyunit code ...

...

Now that I have inserted calls to 'take_screenshot' into the try/except blocks that handle test errors, we now will capture screenshot whenever a test extending this test fails. You can use this test case like this:

import unittest
from wtframework.wtf.testobjects.basetests import WTFBaseTest
from wtframework.wtf.web.webdriver import WTF_WEBDRIVER_MANAGER


class TestScreenCaptureOnFail(ScreenCaptureTestCase):
    """"
    These test cases are expected to fail.  They are here to test 
    the screen capture on failure.
    """

    # Comment out decorator to manually test the screen capture.
    @unittest.expectedFailure
    def test_fail(self):
        driver = WTF_WEBDRIVER_MANAGER.new_driver()
        driver.get('http://www.google.com')
        self.fail()
        #Check your /screenshots folder for a screenshot.

There you go, we have a simple method that automatically captures screenshots upon test failure.


Update: 
I have since then incorporated this into WTFramework,
https://github.com/wiredrive/wtframework

In WTFramework, I extended the basic TestCase class to WatchedTestCase, which allows you to register listeners.
https://github.com/wiredrive/wtframework/blob/master/wtframework/wtf/testobjects/testcase.py

This then allows me to register various listeners like my CaptureScreenShotOnErrorTestWatcher
https://github.com/wiredrive/wtframework/blob/master/wtframework/wtf/testobjects/basetests.py