在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（单点真理/不要重复自己）的精神，我希望将这些位重构成测试方法和方法。

然而，试图这样做却导致了这种丑陋：

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

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

一定有更好的办法。理想情况下，在setUp()/tearDown()中，每个测试方法都没有重复的位（我可以看到在每个方法上重复decorator是如何做到的）。

编辑：将未测试对象视为内部对象，GetResource对象视为第三方对象（我们不会更改）。

我已将GetSlot重命名为GetResource这里，这比上下文管理器是对象进入锁定状态并退出的方式的特定情况更为常见。

Answer 1

如下图所示重写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 2

像您这样调用__enter__和__exit__的问题并不是您已经这样做了：它们可以在with语句之外调用。问题是，如果发生异常，代码没有正确调用对象的__exit__方法的规定。

因此，方法是使用一个decorator将对原始方法的调用包装在一个with语句中。短元类可以透明地将decorator应用于类中名为test*的所有方法-

# -*- coding: utf-8 -*-

from functools import wraps

import unittest

def setup_context(method):
    # the 'wraps' decorator preserves the original function name
    # otherwise unittest would not call it, as its name
    # would not start with 'test'
    @wraps(method)
    def test_wrapper(self, *args, **kw):
        with GetSlot() as slot:
            self._slot = slot
            result = method(self, *args, **kw)
            delattr(self, "_slot")
        return result
    return test_wrapper

class MetaContext(type):
    def __new__(mcs, name, bases, dct):
        for key, value in dct.items():
            if key.startswith("test"):
                dct[key] = setup_context(value)
        return type.__new__(mcs, name, bases, dct)


class GetSlot(object):
    def __enter__(self): 
        return self
    def __exit__(self, *args, **kw):
        print "exiting object"
    def doStuff(self):
        print "doing stuff"
    def doOtherStuff(self):
        raise ValueError

    def getSomething(self):
        return "a value"

def UnderTest(*args):
    return args[0]

class MyTestCase(unittest.TestCase):
  __metaclass__ = MetaContext

  def testFirstThing(self):
      u = UnderTest(self._slot)
      u.doStuff()
      self.assertEqual(u.getSomething(), 'a value')

  def testSecondThing(self):
      u = UnderTest(self._slot)
      u.doOtherStuff()
      self.assertEqual(u.getSomething(), 'a value')

unittest.main()

（我还包括了“GetSlot”的模拟实现，以及示例中的方法和函数，这样我自己就可以测试我建议的decorator和元类）

Answer 3

在不希望with语句在所有资源获取成功时清理内容的情况下操作上下文管理器是^{}设计用来处理的用例之一。

例如（使用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)