Selenium Screenshot on Python unittest Error

This works just fine for failed tests.

class BaseUnitTest(unittest.TestCase):
    @property
    def failureException(self):
        class ScreenshotFailureException(AssertionError):
            def __init__(self_,*args, **kwargs):
                self.driver.save_screenshot('screenshot.png')
                return super(ScreenshotFailureException, self_).__init__(*args, **kwargs)
        ScreenshotFailureException.__name__ = AssertionError.__name__
        return ScreenshotFailureException

I attempted to do something similar by overriding seleniums AbstractEventListener.on_exception() but this gets triggered for every exception that might occur on a page, including any timeout exceptions that occur during a WebDriverWait.

The problem I am having is if I create a decorator (based on above) and apply it to the above class, the decorator only runs for the first instantiation of BaseUnitTest and not for any subsequent test cases. I could apply the decorator to every single test class, but that seems rather tedious.

refer : http://stackoverflow.com/questions/30652804/selenium-screenshot-on-python-unittest-error