在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为GetResourcehere-这比特定情况更通用 - 其中上下文管理器是对象打算进入锁定状态的方式.

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)

这里,如果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()

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

调用__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 5

看来这个讨论10年后仍然有意义！为了添加到@ncoghlan 的优秀答案，看起来从 python 3.11 开始unittest.TestCase通过 helper 方法添加了这个确切的功能！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因为他们有用的固定装置而蜂拥而至）