如何使用类装饰器重写方法而不重新定义类？

Question

在使用 unittest 模块进行单元测试时，如果要用到 App Engine 测试环境，我需要写 setUp 和 tearDown 方法来分别启动和关闭这个测试环境（这里稍微简化了一下）：

class SomeTest(unittest.TestCase):

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

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

  def testSomething(self):
    ...

这样写起来就变得有点麻烦了。我可以写一个基础类 TestCaseWithTestbed，但每次在测试用例中需要自定义 setUp 的时候，我都得记得调用父类的方法。

我觉得用类装饰器来解决这个问题会更优雅。所以我想写成这样：

@WithTestbed
class SomeTest(unittest.TestCase):

  def testSomething(self):
    ...

应用这个装饰器后，测试环境应该可以自动启动。那... 怎么实现这个 WithTestbed 装饰器呢？我现在有以下代码：

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

这个方法在简单情况下是有效的，但有一些严重的问题：

• 测试类的名字变成了 ClsWithTestbed，这个名字会出现在测试输出中。
• 具体的测试类调用 super(SomeTestClass, self).setUp() 时会陷入无限递归，因为 SomeTestClass 现在等于 WithTestbed。

我对 Python 的运行时类型操作有点模糊。那么，怎么才能正确地做到这一点呢？

Answer 1

像这样就可以了：

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):
    ...

Answer 2

这里有一个简单的方法，可以用子类来实现你想要的，而不是使用装饰器：

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"

你只需要在你的测试案例中定义 mySetUp 和 myTearDown，而不是 setUp 和 tearDown。

Answer 3

这个方法看起来有效，解决了问题：

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

有没有人觉得这个方法有什么问题吗？