很多测试都需要在启动的时候做一些事情,然后在结束的时候再把做的事情给清理了。一般的做法是把这些动作写在setUp和tearDown的两个方法里,单元测试框架会负责在开始和结束的时候调用这两个方法。

class SomeTest(unittest.case.TestCase):
    def setUp(self):
        super(SomeTest, self).setUp()
        setup_db()

    def tearDown(self):
        clean_db()
        super(SomeTest, self).tearDown()

这种写法有好几个烦人的地方。首先是Logic Locality不好的问题:setup_db()和clean_db()是分在两处的,中间可能隔着很长一段代码。从视觉上无法直观的指导setup_db()原来和clean_db()是一对的。
其次是很难重用的问题(上纲上线的话就是复杂度不好管理的问题),为了避免重复写公共的setUp和tearDown一般会抽取出一个UsingDbTest这样的基类。这样所有的子类必须记得super(xxx, self).setUp(),否则就会覆盖掉基类的setUp。其次在需要有多个维度的东西需要复用的时候,比如有一个UsingDbTest的基类,有一个UsingNetworkTest的基类,难道让子类继承两个基类么(mixin是不是有点过于复杂了?)。
使用generator可以很好的解决这个问题。首先我们写一个方法来做setUp和tearDown:

@contextlib.contextmanager
def using_db():
    setup_db()
    yield
    clean_db()

这样可以非常清晰地知道setup_db和clean_db是一对的。然后再把这个小的上下文附着到主测试逻辑上:

def apply_context(test, contextmanager):
    contextmanager.__enter__()
    test.addCleanup(lambda: contextmanager.__exit__(None, None, None))

class SomeTest(unittest.case.TestCase):
    def setUp(self):
        apply_context(self, using_db())

这里利用了单元测试的addCleanup的特性,把tearDown转化为回调在setUpd的时候就设置好。利用这种方式,我们可以用组合的方式而不是继承的方式来复用公共的setUp和tearDown的逻辑了。