Python decorator that runs function only if argument test passes

unittest_decorator.py

"""
The `only_if_test_passes` decorator can be used to run a function if and
only if the argument test (unbound `TestCase` method) passes.
"""

import inspect
from unittest import TestCase, TestResult


class TestError(Exception):
    pass


class MyTests(TestCase):

    def test_pass(self):
        # This passes because nothing went wrong
        pass

    def test_fail(self):
        self.fail('This test will always fail')


def only_if_test_passes(test_method):
    # Comments are computed values when passed MyTests.test_pass
    test_case_class = inspect._findclass(test_method)  # MyTests
    test_case_name = test_case_class.__name__  # 'MyTests'
    test_name = test_method.__name__  # 'test_pass'
    # Introspection for Python 2:
    # test_case_class = test_method.im_class
    # test_case_name = test_case_class.__name__  # Same as for Python 3
    # test_name = test_method.if_func.func_name

    def decorator(fn):
        def decorated(*args, **kwargs):
            test_result = TestResult()
            case = test_case_class(test_name)  # MyTests('test_pass')
            case(test_result)
            if test_result.wasSuccessful():
                return fn(*args, **kwargs)
            else:
                raise TestError('Unit test failed: {}.{}'.format(
                    test_case_name, test_name))
        return decorated
    return decorator


@only_if_test_passes(MyTests.test_pass)
def this_will_run():
    print('This should output')


@only_if_test_passes(MyTests.test_fail)
def this_wont_ever_run():
    print("Don't bother; you'll never see this.")


if __name__ == "__main__":
    this_will_run()
    this_wont_ever_run()