In python, is there a good idiom for using context managers in setup/teardown

Question

I am finding that I am using plenty of context managers in Python. However, I have been testing a number of things using them, and I am often needing the following:

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

When this gets to many tests, this is clearly going to get boring, so in the spirit of SPOT/DRY (single point of truth/dont repeat yourself), I'd want to refactor those bits into the test setUp() and tearDown() methods.

However, trying to do that has lead to this ugliness:

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

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

There must be a better way to do this. Ideally, in the setUp()/tearDown() without repetitive bits for each test method (I can see how repeating a decorator on each method could do it).

Edit: Consider the undertest object to be internal, and the GetResource object to be a third party thing (which we aren't changing).

I've renamed GetSlot to GetResource here—this is more general than specific case—where context managers are the way which the object is intended to go into a locked state and out.

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

The problem with calling __enter__ and __exit__ as you did, is not that you have done so: they can be called outside of a with statement. The problem is that your code has no provision to call the object's __exit__ method properly if an exception occurs.

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.

Answer 6

Looks like this discussion is still relevant 10 years later! To add to @ncoghlan's excellent answer it looks like unittest.TestCase added this exact functionality via the enterContext helper method as of python 3.11! From the docs:

enterContext(cm)

Enter the supplied context manager. If successful, also add its __exit__() method as a cleanup function by addCleanup() and return the result of the __enter__() method.

New in version 3.11.

It looks like this precludes the need to manually addCleanup() to close the stack of context managers, as it's added when you provide the context manager to enterContext. So it seems like all that's needed nowadays is:

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

(I guess unittest got tired of everyone flocking to pytest because of their helpful fixtures)