Automatic screenshots when test fail by Selenium Webdriver in Python

Here is a solution using a decorator that wrapps every method on a class that starts test_ with a wrapper that takes a screenshot if the method raises and Exception. The browser_attr is used to tell the decorator how to obtain the web browser (driver).

from functools import partialmethod


def sreenshotOnFail(browser_attr='browser'):
    def decorator(cls):
        def with_screen_shot(self, fn, *args, **kwargs):
            """Take a Screen-shot of the drive page, when a function fails."""
            try:
                return fn(self, *args, **kwargs)
            except Exception:
                # This will only be reached if the test fails
                browser = getattr(self, browser_attr)
                filename = 'screenshot-%s.png' % fn.__name__
                browser.get_screenshot_as_file(filename)
                print('Screenshot saved as %s' % filename)
                raise

        for attr, fn in cls.__dict__.items():
            if attr[:5] == 'test_' and callable(fn):
                setattr(cls, attr, partialmethod(with_screen_shot, fn))

        return cls
    return decorator


@sreenshotOnFail()
class TestDemo(unittest.TestCase):
    def setUp(self):
        """Set up the Firefox Browser and the Tear Down."""
        self.browser = webdriver.Firefox()

    def test_demo2(self):
        """A test case that fails because assert."""
        self.driver.get("https://stackoverflow.com")
        self.assertEqual(True, False)

You can start form exploring self._outcome.errors[1] where it is possible to find information about the errors.
i.e. below code will work only for assertion errors

 def tearDown(self):
        if self._outcome.errors[1][1] and hasattr(self._outcome.errors[1][1][1], 'actual'):
            self.driver.save_screenshot(self._testMethodName + '.png')

try:
   page2.fill_up_company_info()
except Exception as e:
   with pytest.raises(WebDriverException):
            pytest.fail( f"{e} for error message",       
allure.attach(self.driver.get_screenshot_as_png(),name="screenshot",attachment_type=AttachmentType.PNG))

For Django 2.2.2 (which uses unittest) in my class for selenium tests, which inherited from StaticLiveServerTestCase I've overrided _feedErrorsToResult method. Also this approach provides tricky way to know name of called method for convinient screenshot investigation.

@classmethod
def _feedErrorsToResult(cls, result, errors):
    """
    Overriding private method at %library root%/Lib/unittest/case.py
    so you can take screenshot with any failed test and find name of the method
    """
    if SELENIUM_TAKE_SCREENSHOTS:
        for test, exc_info in errors:
            if exc_info is not None:
                now = datetime.now().strftime('%y-%m-%d_%H-%M-%S')
                test_name = exc_info[2].tb_frame.f_locals["test_case"]._testMethodName
                # noinspection PyUnresolvedReferences
                cls.selenium.get_screenshot_as_file('%s/%s-%s-%s.png' % (SELENIUM_SCREENSHOTS_PATH, cls.__name__, test_name, now))
    # noinspection PyUnresolvedReferences
    super()._feedErrorsToResult(cls, result, errors)