python - 在 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()没有重复位（我可以看到在每个方法上重复装饰器是如何做到的）。

编辑：考虑undertest对象是内部的，GetResource对象是第三方的东西（我们没有改变）。

我在这里重命名GetSlot了GetResource——这比特定情况更普遍——上下文管理器是对象进入锁定状态和退出的方式。

Answer 1

How about overriding unittest.TestCase.run() as illustrated below? This approach doesn't require calling any private methods or doing something to every method, which is what the questioner wanted.

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()

This approach also allows passing the TestCase instance to the context manager, if you want to modify the TestCase instance there.

Answer 2

Manipulating context managers in situations where you don't want a with statement to clean things up if all your resource acquisitions succeed is one of the use cases that contextlib.ExitStack() is designed to handle.

For example (using addCleanup() rather than a custom tearDown() implementation):

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

That's the most robust approach, since it correctly handles acquisition of multiple resources:

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)

Here, if GetOtherResource() fails, the first resource will be cleaned up immediately by the with statement, while if it succeeds, the pop_all() call will postpone the cleanup until the registered cleanup function runs.

If you know you're only ever going to have one resource to manage, you can skip the with statement:

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

However, that's a bit more error prone, since if you add more resources to the stack without first switching to the with statement based version, successfully allocated resources may not get cleaned up promptly if later resource acquisitions fail.

You can also write something comparable using a custom tearDown() implementation by saving a reference to the resource stack on the test case:

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()

Alternatively, you can also define a custom cleanup function that accesses the resource via a closure reference, avoiding the need to store any extra state on the test case purely for cleanup purposes:

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 fixtures are very close to your idea/style, and allow for exactly what you want:

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__问题是如果发生异常，您的代码无法正确调用对象的方法。

So, the way to do it is to have a decorator that will wrap the call to your original method in a withstatement. A short metaclass can apply the decorator transparently to all methods named test* in the class -

# -*- 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()

(I also included mock implementations of "GetSlot" and the methods and functions in your example so that I myself could test the decorator and metaclass I am suggesting on this answer)

Answer 5

I'd argue you should separate your test of the context manager from your test of the Slot class. You could even use a mock object simulating the initialize/finalize interface of slot to test the context manager object, and then test your slot object separately.

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()

This makes code simpler, prevents concern mixing and allows you to reuse the context manager without having to code it in many places.