在Python中，是否有一种好的惯用语用于在设置/拆卸中使用上下文管理器？

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，这是一个更通用的名称，其中上下文管理器是使该对象进入锁定状态和退出锁定状态的方式。

- Danny Staple

---

我不明白你的setUp/tearDown方法有什么问题，对我来说看起来完全没问题。我想另一种选择是创建一个使用with语句的装饰器，并自动将其应用于所有方法，但这会增加更多的工作量而没有实际好处。 - interjay

我认为'__'方法是私有的和"神奇"的方法，不应该被显式调用。然而，考虑到这是在测试环境中，也许这已经足够了。 - Danny Staple

设置和拆卸是两者中更清洁的。我认为GetSlot应该有适当的API，可以在没有上下文管理器的情况下使用。你苦于找到最干净的方法来做这件事实证了GetSlot需要改进。除非GetSlot不是你的代码，否则我收回所有话。 - Chris Lacasse

采用你现有的解决方案，完全可以在测试用例中调用“魔术”方法。 - Ferdinand Beyer

如果代码不是我的代码，我可以将其封装，以便拥有一个干净的外部API，尽管也许我准备接受@FerdinandBeyer的回答作为合适的答案。将它们添加为答案（我可以给他们信用）。 - Danny Staple

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实例传递给上下文管理器。

- cjerdonek

---

我喜欢这个 - 这是一个简单的解决方案。 - Danny Staple

我认为答案不够清晰。更好地解释一下test()方法是在上下文中执行的，并且在断言之前可以执行doStuff()或doOtherStuff()方法。 - Alan Franzoni

我在尝试理解http://flask.pocoo.org/snippets/26/时遇到了这个问题，似乎在那里使用`__call__`类似于`run`。我想知道覆盖`__call__`与`run`之间的区别是什么？ - BenjaminGolder

@BenjaminGolder TestCase.__call__ 被定义为 return self.run(*args, **kwds)，因此区别似乎只是风格上的。 - Brian M. Hunt

顺序完全错了。如果你从许多具有setUp()/tearDown()和现在还有run()方法的mixin中派生MyTest，那么它首先执行所有的run()，然后是所有的setUp()，最后是所有的tearDown()。但我们需要上下文跨越属于特定mixin的setUp()和tearDown()之间。 - Velkan

Answer 2

contextlib.ExitStack()是为了处理在情况中操作上下文管理器，当您不希望使用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)

- ncoghlan

---

+1 这绝对比覆盖“run”方法更好。唯一的缺点是，上下文管理器无法在这种情况下处理异常。 - Neil G

最后一个 tearDown 的例子有点危险，因为如果 setUp 在 with 语句之后失败，tearDown 将不会被调用。建议使用 addCleanup。 - user1338062

一个带有本地函数的例子将会在这里非常有帮助，例如：def cleaner(): stack.pop_all().close() self.addCleanup(cleaner) - Bob Stein

@BobStein 我不确定我理解了，因为那段代码等同于第一个例子，只是多了一个包装函数作为回调，而不是直接传入绑定的 close 方法。如果你喜欢，当然可以这样做，但这会稍微慢一些，没有实际好处。 - ncoghlan

不是在那段代码中，你说得对。但是许多清理工作需要多个函数调用。（我的情况需要2次调用，1次条件语句。）因此，一个包装器将展示这种技术的更多可扩展性。此外，将函数视为对象已经是一种奇怪的想法了。也许最好还是显式一些。我现在要编辑您的答案，加上一个示例包装器，以便澄清。请随意恢复或编辑，视情况而定。我从您的addCleanup()点中获得了很多收获，但也许这与您的ExitStack()点无关。 - Bob Stein

@BobStein 噢，现在我明白了。是的，在闭包中存储状态可以成为存储在测试用例上的不错替代品。 - ncoghlan

Answer 3

pytest 的固件非常接近于您的想法/风格，并且可完全满足您的要求：

import pytest
from code.to.test import foo

@pytest.fixture(...)
def resource():
    with your_context_manager as r:
        yield r

def test_foo(resource):
    assert foo(resource).bar() == 42

- Dima Tisnek

---

简单的解决方案：使用pytest - derchambers

Answer 4

调用 __enter__ 和 __exit__ 的问题不在于你这样做了：它们可以在 with 语句之外被调用。问题在于，如果发生异常，你的代码没有正确调用对象的 __exit__ 方法。

因此，正确的方法是使用一个装饰器来将原始方法的调用包装在 with 语句中。一个简短的元类可以透明地将装饰器应用于类中所有以 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”以及你的示例中所涉及的方法和函数的模拟实现，这样我就可以测试我在回答中建议的装饰器和元类。

- jsbueno

---

参考 Python 文档：http://docs.python.org/library/unittest.html#unittest.TestCase.tearDown，tearDown 方法会在每次测试结束后被调用，即使有失败的测试用例和异常情况（setUp 方法则在每次测试前执行）。MetaClass/MetaContext 肯定只会在整个测试用例 MyTestCase 运行时进入/退出，而不是在单独的单元测试中。这可能意味着测试之间存在交互。 - Danny Staple

@DannyStaple：至于我的代码：我将每个单独的测试方法都包装在一个装饰器中，该装饰器调用with语句内的文本 - 在每次测试执行时都会运行enter和exit。在test_wrapper函数中插入一些打印语句，亲自看看吧。至于你的原始代码，很高兴知道.__exit__将被调用。从技术上讲，.__exit__应该传递有关发生的异常的信息，但我同意在大多数情况下这不应该是问题。 - jsbueno

啊，我明白了 - 你正在访问对象字典并将上下文应用于每个以“test”开头的项目，这将是测试方法。@wraps装饰器很方便 - 我之前在包装时遇到过丢失方法名称的问题。 - Danny Staple

Answer 5

看起来这个讨论十年后仍然相关！除了@ncoghlan的优秀回答，根据Python 3.11的文档，unittest.TestCase通过enterContext辅助方法添加了这个确切的功能！

enterContext(cm)

进入提供的上下文管理器。如果成功，还将其__exit__()方法作为清理函数添加到addCleanup()中，并返回__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

- BentHam

Answer 6

我认为你应该将上下文管理器的测试与 Slot 类的测试分开。您甚至可以使用模拟 slot 初始化/结束接口的 mock 对象来测试上下文管理器对象，然后分别测试 slot 对象。

from unittest import TestCase, main

class MockSlot(object):
    initialized = False
    ok_called = False
    error_called = False

    def initialize(self):
        self.initialized = True

    def finalize_ok(self):
        self.ok_called = True

    def finalize_error(self):
        self.error_called = True

class GetSlot(object):
    def __init__(self, slot_factory=MockSlot):
        self.slot_factory = slot_factory

    def __enter__(self):
        s = self.s = self.slot_factory()
        s.initialize()
        return s

    def __exit__(self, type, value, traceback):
        if type is None:
            self.s.finalize_ok()
        else:
            self.s.finalize_error()


class TestContextManager(TestCase):
    def test_getslot_calls_initialize(self):
        g = GetSlot()
        with g as slot:
            pass
        self.assertTrue(g.s.initialized)

    def test_getslot_calls_finalize_ok_if_operation_successful(self):
        g = GetSlot()
        with g as slot:
            pass
        self.assertTrue(g.s.ok_called)

    def test_getslot_calls_finalize_error_if_operation_unsuccessful(self):
        g = GetSlot()
        try:
            with g as slot:
                raise ValueError
        except:
            pass

        self.assertTrue(g.s.error_called)

if __name__ == "__main__":
    main()

这样可以使代码更简单，避免关注点混杂，并且可以在多个地方重复使用上下文管理器，而无需在许多地方编写它。

- Alan Franzoni

---

我认为模拟可能是解决方案，所以我会给你一个+1。原始代码（在许多情况下 - 我现在已经将问题概括了一点）允许您构建某些内容，但在进入上下文之前它并没有被锁定。请注意，在原始问题中，上下文和插槽都不是SUT（受测主体） - 它们是资源。 - Danny Staple

现在回过头来看，我认为模拟资源正是我现在要做的事情，因为根据你所指出的，上面的代码中，资源根本没有被测试到。您可能希望通过使用模拟来避免资源可能产生的任何副作用。我只会模拟SUT所需的资源接口，甚至可以创建助手来构建模拟资源（如果我需要它）。 - Danny Staple

这个答案仍然没有解决如何以DRY的方式将上下文管理器应用于TestCase类中的测试的问题。 - cjerdonek

cjerdonek：setUp/tearDown 本身就是上下文管理器的行为，尽管它在语法风格上与 Python 不同。如果你一定要使用上下文管理器，那么这可能是一个 API 的问题，你应该将上下文管理器与实际代码解耦。顺便说一句，Danny 提出的 setUp/tearDown 解决方案虽然不太美观，但应该可以百分之百地工作。 - Alan Franzoni

Alan，我知道Danny的代码是有效的，但他的“解决方案”是他问题的一部分。他正在寻求一种更少丑陋的方式来以DRY的方式将上下文管理器应用于TestCase类中的测试（例如，不需要访问上下文管理器的私有方法）。也许我的评论应该更简单地表明这个答案仍然没有回答他最初的问题。我提供了一种方法来做到这一点，允许使用上下文管理器本身。 - cjerdonek