pythonist pythonist - 6 months ago 24
Python Question

Perform a screenshot upon an Exception

Hey is there a way to capture a screenshot upon a an Exception, any Exception. My 'failed' solution which is placed in the

BaseTestCase
,
unittest.TestCase
subclass:

class BaseTestCase(unittest.TestCase):
@classmethod
def setUpClass(cls):
"""some code"""

@property
def failureException(self):
class MyFailureException(Exception):
def __init__(self_, *args, **kwargs):
screenshot_dir = '../error_log'
self.driver.save_screenshot('{0}/{1}.jpeg'.format(screenshot_dir, self.id()))
return super().__init__(*args, **kwargs)

MyFailureException.__name__ = Exception.__name__
return MyFailureException


Originally it was
AssertionError
instead of the
Exception
but it would only catch assertion errors and I'm more interested in other kinds of errors

Answer

To take a screenshot when an error or failure occurs, check if there is an exception currently being handled in the tearDown method:

import unittest, sys, exceptions
from selenium import webdriver


class TestCaseBase(unittest.TestCase):

    def setUp(self):
        self.driver = webdriver.Firefox()

    def tearDown(self):
        type, value, traceback = sys.exc_info()
        if type is exceptions.AssertionError :
            self.driver.save_screenshot(r'screenshot-failure.png')
        elif type is exceptions.Exception :
            self.driver.save_screenshot(r'screenshot-error.png')

        self.driver.quit()

class MyTestCase(TestCaseBase):

    def test_should_take_screenshot_on_failure(self):
        self.driver.get("http://stackoverflow.com/")
        self.assertTrue(False)

    def test_should_take_screenshot_on_error(self):
        self.driver.get("http://stackoverflow.com/")
        raise Exception("my exception")


if __name__ == '__main__':
    unittest.main()

You could also override the methods TestResult.addError and TestResult.addFailure:

import unittest
from selenium import webdriver


class TestCaseBase(unittest.TestCase):

    def setUp(self):
        self.driver = webdriver.Firefox()

    def tearDown(self):
        self.driver.quit()

    def run(self, result=None):
        super(TestCaseBase, self).run(TestResultEx(result, self))


class MyTestCase(TestCaseBase):

    def test_should_take_screenshot_on_failure(self):
        self.driver.get("http://stackoverflow.com/")
        self.assertTrue(False)

    def test_should_take_screenshot_on_error(self):
        self.driver.get("http://stackoverflow.com/")
        raise Exception("my exception")


class TestResultEx(object):

    def __init__(self, result, testcase):
        self.result = result
        self.testcase = testcase

    def __getattr__(self, name):
        return object.__getattribute__(self.result, name)

    def addError(self, test, err):
        self.result.addError(test, err)
        self.testcase.driver.save_screenshot(r'screenshot-error.png')

    def addFailure(self, test, err):
        self.result.addFailure(test, err)
        self.testcase.driver.save_screenshot(r'screenshot-failure.png')


if __name__ == '__main__':
    unittest.main()