pythonist pythonist - 1 year ago 61
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 Source

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()