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.