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

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)

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.

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)

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

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)