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)