pythonunit-testingcontextmanager

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


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.


Solution

  • 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 = self.enterContext(GetResource()) # if you need a reference to it in tests
        # self._resource implicitly released during cleanups after tearDown()
    

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