Python: unittest setUp and tearDown with a ContextManager

Python unittest follows the jUnit structure, but is extremely awkward.  One of the more awkward portions are the use of setUp and tearDown methods.  Python has an elegant way of handling setup and teardown, it's called a ContextManager.  So let's add it.

import unittest
from functools import wraps
from contextlib import contextmanager

def addContextHandler(fn, ctx):
    @wraps(fn)
    def helper(self, *a, **kw):
        if not hasattr(self, ctx):
            return fn(self, *a, **kw)

        with getattr(self, ctx)():
            return fn(self, *a, **kw)

    return helper

unittest.TestCase.run = addContextHandler(unittest.TestCase.run, 'contextTest')

class TestOne(unittest.TestCase):

    @contextmanager
    def contextTest(self):
        print "starting context"
        yield
        print "ending context"

    def testA(self):
        print "testA"
        self.assertTrue(True)

    def testB(self):
        print "testB"
        self.assertTrue(True)

    def testC(self):
        print "testC"
        self.assertTrue(True)
        
if __name__ == "__main__":
    unittest.main()

Or if you would rather not do a dynamic patch you can subclass it.

import unittest

class MyTestCase(unittest.TestCase):
    ctx = 'contextTest'

    def run(self, *a, **kw):
        if not hasattr(self, self.ctx):
            return super(MyTestCase, self).run(*a, **kw)

        with getattr(self, self.ctx)():
            return super(MyTestCase, self).run(*a, **kw)
     

§