在Python中，有什么好的习惯用法用于setup/teardown的上下文管理器吗

Question

我发现自己在Python中使用了很多上下文管理器。不过，我在测试一些东西时，常常需要以下这些：

class MyTestCase(unittest.TestCase):
  def testFirstThing(self):
    with GetResource() as resource:
      u = UnderTest(resource)
      u.doStuff()
      self.assertEqual(u.getSomething(), 'a value')

  def testSecondThing(self):
    with GetResource() as resource:
      u = UnderTest(resource)
      u.doOtherStuff()
      self.assertEqual(u.getSomething(), 'a value')

当测试的数量变多时，这样做显然会变得很无聊，所以为了遵循SPOT/DRY（单一真相点/不要重复自己）的原则，我想把这些部分重构到测试的setUp()和tearDown()方法里。

然而，尝试这样做却导致了以下的麻烦：

  def setUp(self):
    self._resource = GetSlot()
    self._resource.__enter__()

  def tearDown(self):
    self._resource.__exit__(None, None, None)

肯定有更好的方法来处理这个问题。理想情况下，希望在setUp()和tearDown()中，不用在每个测试方法里重复相同的代码（我能理解在每个方法上重复使用装饰器可以解决这个问题）。

编辑：可以把正在测试的对象看作是内部的，而GetResource对象则是一个第三方的东西（我们不打算去修改它）。

我在这里把GetSlot改名为GetResource——这个名字更通用，而不是特定的案例——上下文管理器是这个对象进入和退出锁定状态的方式。

Answer 1

在一些情况下，你可能不想让 with 语句在所有资源都成功获取时自动清理资源，这时就可以用到 contextlib.ExitStack()。

举个例子（这里用 addCleanup() 而不是自定义的 tearDown() 实现）：

def setUp(self):
    with contextlib.ExitStack() as stack:
        self._resource = stack.enter_context(GetResource())
        self.addCleanup(stack.pop_all().close)

这种方法是最稳妥的，因为它能正确处理多个资源的获取：

def setUp(self):
    with contextlib.ExitStack() as stack:
        self._resource1 = stack.enter_context(GetResource())
        self._resource2 = stack.enter_context(GetOtherResource())
        self.addCleanup(stack.pop_all().close)

在这个例子中，如果 GetOtherResource() 失败了，第一个资源会立刻被 with 语句清理掉；而如果成功了，pop_all() 的调用会把清理工作推迟到注册的清理函数运行时再执行。

如果你知道自己只需要管理一个资源，可以省略 with 语句：

def setUp(self):
    stack = contextlib.ExitStack()
    self._resource = stack.enter_context(GetResource())
    self.addCleanup(stack.close)

不过，这样做有点容易出错，因为如果你在不先切换到基于 with 语句的版本的情况下添加更多资源，如果后续获取资源失败，已经成功分配的资源可能不会及时被清理。

你也可以通过在测试用例中保存资源堆栈的引用，使用自定义的 tearDown() 实现来写类似的代码：

def setUp(self):
    with contextlib.ExitStack() as stack:
        self._resource1 = stack.enter_context(GetResource())
        self._resource2 = stack.enter_context(GetOtherResource())
        self._resource_stack = stack.pop_all()

def tearDown(self):
    self._resource_stack.close()

另外，你还可以定义一个自定义的清理函数，通过闭包引用来访问资源，这样就不需要在测试用例中额外存储状态仅仅为了清理：

def setUp(self):
    with contextlib.ExitStack() as stack:
        resource = stack.enter_context(GetResource())

        def cleanup():
            if necessary:
                one_last_chance_to_use(resource)
            stack.pop_all().close()

        self.addCleanup(cleanup)

Answer 2

那我们可以试试重写一下 unittest.TestCase.run()，就像下面这样做。这种方法不需要调用任何私有方法，也不需要对每个方法做什么，这正是提问者想要的。

from contextlib import contextmanager
import unittest

@contextmanager
def resource_manager():
    yield 'foo'

class MyTest(unittest.TestCase):

    def run(self, result=None):
        with resource_manager() as resource:
            self.resource = resource
            super(MyTest, self).run(result)

    def test(self):
        self.assertEqual('foo', self.resource)

unittest.main()

这种方法还允许将 TestCase 实例传递给上下文管理器，如果你想在那儿修改 TestCase 实例的话。

Answer 3

看起来这个讨论在10年后依然很有意义！为了补充一下@ncoghlan的精彩回答，从Python 3.11开始，unittest.TestCase增加了一个非常实用的功能，通过enterContext这个辅助方法来实现！在官方文档中这样写道：

enterContext(cm)

进入提供的上下文管理器。如果成功，还会通过addCleanup()将它的__exit__()方法添加为清理函数，并返回__enter__()方法的结果。

这是在3.11版本中新增加的。

这意味着你不再需要手动调用addCleanup()来关闭上下文管理器的堆栈，因为当你把上下文管理器传给enterContext时，它会自动添加。所以现在只需要做的就是：

def setUp(self):
    self._resource = GetResource() # if you need a reference to it in tests
    self.enterContext(GetResource())
    # self._resource implicitly released during cleanups after tearDown()

（我想unittest是厌倦了大家都去用pytest，因为它的帮助功能太好用了）