Using a class decorator, how to override a method without redefining the class?

Question

For unit tests (using the unittest module) that use the App Engine testbed, I need setUp and tearDown methods to activate and deactivate the testbed, respectively (slightly simplified):

class SomeTest(unittest.TestCase):

  def setUp(self):
    self.testbed = testbed.Testbed()
    self.testbed.activate()

  def tearDown(self):
    self.testbed.deactivate()

  def testSomething(self):
    ...

This quickly becomes a burden to write. I could write a base class TestCaseWithTestbed, but then I'd have to remember to call the superclass method each time I need a custom setUp in one of the test cases.

I thought it would be more elegant to solve this with a class decorator instead. So I'd like to write:

@WithTestbed
class SomeTest(unittest.TestCase):

  def testSomething(self):
    ...

With this decorator applied, the testbed should just be activated magically. So... how to implement the WithTestbed decorator? I currently have the following:

def WithTestbed(cls):
  class ClsWithTestbed(cls):

    def setUp(self):
      self.testbed = testbed.Testbed()
      self.testbed.activate()
      cls.setUp(self)

    def tearDown(self):
      cls.tearDown(self)
      self.testbed.deactivate()

  return ClsWithTestbed

This works for simple cases, but has some serious problems:

• The name of the test class becomes ClsWithTestbed and this shows up in the test output.
• Concrete test classes calling super(SomeTestClass, self).setUp() end up in an infinite recursion, because SomeTestClass is now equal to WithTestbed.

I'm a bit hazy on Python's runtime type manipulation. So, how to do this the Right Way?

Answer 1

This appears to work and solve the problems:

def WithTestbed(cls):
  def DoNothing(self):
    pass

  orig_setUp = getattr(cls, 'setUp', DoNothing)
  orig_tearDown = getattr(cls, 'tearDown', DoNothing)

  def setUp(self):
    self.testbed = testbed.Testbed()
    self.testbed.activate()
    orig_setUp(self)
  def tearDown(self):
    orig_tearDown(self)
    self.testbed.deactivate()

  cls.setUp = setUp
  cls.tearDown = tearDown
  return cls

Does anyone see any problems with this approach?

Answer 2

Here's a simple way to do what you're asking with subclassing instead of a decorator:

class TestCaseWithTestBed(unittest.TestCase):

  def setUp(self):
    self.testbed = testbed.Testbed()
    self.testbed.activate()
    self.mySetUp()

  def tearDown(self):
    self.myTearDown()
    self.testbed.deactivate()

  def mySetUp(self): pass
  def myTearDown(self): pass

class SomeTest(TestCaseWithTestBed):
  def mySetUp(self):
    "Insert custom setup here"

All you have to do is define mySetUp and myTearDown in your test cases instead of setUp and tearDown.

Answer 3

Something like this would work:

def WithTestbed(cls):
    cls._post_testbed_setUp = getattr(cls, 'setUp', lambda self : None)
    cls._post_testbed_tearDown = getattr(cls, 'tearDown', lambda self : None)

    def setUp(self):
        self.testbed = testbed.Testbed()
        self.testbed.activate()
        self._post_testbed_setUp()

    def tearDown(self):
        self.testbed.deactivate()
        self._post_testbed_tearDown()

    cls.setUp = setUp
    cls.tearDown = tearDown
    return cls

@WithTestbed
class SomeTest(object):
    ...