How to execute code only on test failures with python unittest2?

Question

I have some class-based unit tests running in python's unittest2 framework. We're using Selenium WebDriver, which has a convenient save_screenshot() method. I'd like to grab a screenshot in tearDown() for every test failure, to reduce the time spent debugging why a test failed.

However, I can't find any way to run code on test failures only. tearDown() is called regardless of whether the test succeeds, and I don't want to clutter our filesystem with hundreds of browser screenshots for tests that succeeded.

How would you approach this?

Answer 1

Found a solution - I can override failureException:

@property
def failureException(self):
    class MyFailureException(AssertionError):
        def __init__(self_, *args, **kwargs):
            self.b.save_screenshot('%s.png' % self.id())
            return super(MyFailureException, self_).__init__(*args, **kwargs)
    MyFailureException.__name__ = AssertionError.__name__
    return MyFailureException

This seems incredibly hacky but it seems to work so far.

Answer 2

Here is similar approach to @craigds answer, but with directory support and better compatibility with Python 3:

@property
def failureException(self):
    class MyFailureException(AssertionError):
        def __init__(self_, *args, **kwargs):
            screenshot_dir = 'reports/screenshots'
            if not os.path.exists(screenshot_dir):
                os.makedirs(screenshot_dir)
            self.driver.save_screenshot('{0}/{1}.png'.format(screenshot_dir, self.id()))
            return super(MyFailureException, self_).__init__(*args, **kwargs)
    MyFailureException.__name__ = AssertionError.__name__
    return MyFailureException

This was actually found in this blog.

I've extended it further more with argparse:

parser.add_argument("-r", "--reports-dir", action="store",   dest="dir",      help="Directory to save screenshots.", default="reports")     

so the dir can be specified dynamically either by system variable or passed argument:

screenshot_dir = os.environ.get('REPORTS_DIR', self.args.dir) + '/screenshots'

This is especially useful, if you've additional wrapper to run all your scripts, like a base class.