Execute multiple tests in the same context managers

Question

I have a test setup like this:

def test1():
    with Manager1() as m1, m1.get_m2() as m2:
        do_test1(m1, m2)

def test2():
    with Manager1() as m1, m1.get_m2() as m2:
        do_test2(m1, m2)

which I run with nosetests. Now it turns out that the context managers' __enter__() and __exit__() methods are expensive, so I would like to change to a setup like:

def test():
    with Manager1() as m1, m1.get_m2() as m2:
        do_test1(m1, m2)
        do_test2(m1, m2)

but I would like nose to still report separately for each test. I looked into module level and class level setup() and teardown() methods, but they don't seem to provide the same guarantees for cleanup in the case of exceptions. Is there a clean way to get nose to collaborate with context managers?

Answer 1

If you are worried about an exception being raised in the setup causing the teardown method not to be called then you could use the ExitStack class in contextlib. The ExitStack can make sure Your manager classes are closed properly if the case of an exception in setup.

from io import StringIO
from contextlib import ExitStack

managers = []
close_managers = None

def setup():
    global close_managers
    with ExitStack() as exit_stack:
        managers[:] = [exit_stack.enter_context(StringIO()) for _ in range(2)]
        # any exception raised before this line will will cause all the __exit__
        # methods of the given context managers to be called. So this should
        # be the last part of setup
        close_managers = exit_stack.pop_all().close

def teardown():
    m1, m2 = managers
    assert m1.getvalue() == 'test1 m1\ntest2 m1\n'
    assert m2.getvalue() == 'test1 m2\ntest2 m2\n'
    close_managers()

def test1():
    m1, m2 = managers
    m1.write('test1 m1\n')
    m2.write('test1 m2\n')    

def test2():
    m1, m2 = managers
    m1.write('test2 m1\n')
    m2.write('test2 m2\n')