python 'with' statement and its use in a class

Question

I have a TestCase which uses sockets for testing purposes. As you should always close a socket even in the case of some error, I created a context-manager class abstracting a socket.

A test fixture would look like the following, where MyClassToTest is the actual class to test which uses the socket device internally.

with SocketContextManager() as device:
    foo = MyClassToTest(device)
    # make tests with foo

I would like to avoid having these two line repeated in each of the Test fixture, but placed consistently into setUp. But how should I do this? The following code

def setUp(self):
    with SocketContextManager() as device:
        self.foo = MyClassToTest(device)

does not work as the device will be closed at the end of setUp. Is there a way to handle the instantiation of the context manager like this, or do I have to repeat it in every test fixture?

Answer 1

According to the documentation for tearDown:

This is called even if the test method raised an exception

So you can just open the socket in setUp, and close it in tearDown. Even if your test case raises an exception, the socket will still be closed.

Answer 2

This is a very interesting question. As BrenBarn pointed out, the unittest framework doesn't support doing what you want, but it seems to me that there's no particular reason why you couldn't adapt it to fit. The setUp/tearDown pairing is a hangover from other languages that didn't have generators.

The code below defines a new 'ContextTest' class which amalgamates both setUp and tearDown methods into a single generator that both builds and destroys the context for the test. You can drop your with statement into the context() method along with any other boilerplate.

#!/usr/bin/python3.3
import unittest

class ContextTest(unittest.TestCase):
    """A unit test where setUp/tearDown are amalgamated into a
    single generator"""
    def context(self):
        """Put both setUp and tearDown code in this generator method
        with a single `yield` between"""
        yield

    def setUp(self):
        self.__context = self.context()
        next(self.__context)
    def tearDown(self):
        for _ in self.__context:
            raise RuntimeError("context method should only yield once")

from contextlib import closing
from urllib.request import urlopen

class MyTest(ContextTest):
    def context(self):
        with closing(urlopen('http://www.python.org')) as self.page:
            yield

    def testFoo(self):
        self.assertIsNotNone(self.page)

if __name__=='__main__':
    unittest.main()