在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（这比具体情况更普遍），其中上下文管理器是对象进入锁定状态和退出锁定状态的方式。

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

如果所有资源获取都成功，则在您不希望 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)

这里，如果 < code>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 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

Answer 4

看来这个讨论10年后仍然有意义！要添加到 @ncoghlan 的出色答案，看起来 unittest.TestCase 通过从 python 3.11 开始，enterContext 辅助方法！来自 文档：

进入上下文（cm）

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

3.11 版本中的新增功能。

看起来这不需要手动 addCleanup() 来关闭上下文管理器堆栈，因为它是在您向 enterContext 提供上下文管理器时添加的。所以看来现在所需要的就是：（

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

我猜 unittest 已经厌倦了每个人都涌向 pytest 因为他们有用的装置）

Answer 5

像您一样调用 __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”的模拟实现以及示例中的方法和函数，以便我自己可以测试我建议的装饰器和元类这个答案）

Answer 6

我认为您应该将上下文管理器的测试与 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()

这使得代码更简单，防止关注点混合，并允许您重用上下文管理器，而无需在许多地方对其进行编码。