CODEWITHSUNDEEP
pythonpython-unittestwith-statement
kriss
kriss

Reputation: 24207

Replacing python with_statement by setUp and tearDown in Unittest

In a test suite I have some code organized as below, context is some persistent object that is deleted when exiting the with block:

class Test(TestCase):
    def test_action1(self):
        with create_context() as context:
            context.prepare_context()
            context.action1()
            self.assertTrue(context.check1())
            
    def test_action2(self):
        with create_context() as context:
            context.prepare_context()
            context.action2()
            self.assertTrue(context.check2())

It's obvious that code has some repetition of setup boilerplate in both tests, hence I would like to use setUp() and tearDown() methods to factorize that boilerplate.

But I don't know how to extract the with_statement. What I came up with is something like this:

class Test(TestCase):
    def setUp(self):
        self.context = create_context()
        self.context.prepare_context()

    def tearDown(self):
        del self.context()

    def test_action1(self):
        self.context.action1()
        self.assertTrue(self.context.check1())
            
    def test_action2(self):
        self.context.action2()
        self.assertTrue(self.context.check2())

But I believe this is not really equivalent when test fails, also having to put an explicit delete in tearDown() doesn't feel right.

What is the correct way to change my with_statement code to setUp() and tearDown() style ?

Upvotes: 2

Views: 287

Answers (2)

Tom Wojcik
Tom Wojcik

Reputation: 6189

You might want to use contextlib.ExitStack for that.

A context manager that is designed to make it easy to programmatically combine other context managers and cleanup functions, especially those that are optional or otherwise driven by input data.

import contextlib
from unittest import TestCase


class Test(TestCase):
    def setUp(self) -> None:
        stack = contextlib.ExitStack()
        self.context = stack.enter_context(create_context())  # create_context is your context manager
        self.addCleanup(stack.close)

    def test_action1(self):
        self.context.prepare_context()
        self.context.action1()
        self.assertTrue(self.context.check1())

or if you want to have some control over the teardown or use multiple context managers this will be better

import contextlib
from unittest import TestCase


class Test(TestCase):
    def setUp(self):
        with contextlib.ExitStack() as stack:
            self.context = stack.enter_context(create_context())  # create_context is your context manager
            self._resource_stack = stack.pop_all()

    def tearDown(self):
        self._resource_stack.close()

    def test_action1(self):
        self.context.prepare_context()
        self.context.action1()
        self.assertTrue(self.context.check1())

Upvotes: 2

jrmylow
jrmylow

Reputation: 689

I'm not 100% sure about the setUp() and tearDown(), but the methods defined in context managers __enter__ and __exit__ sound like they do what you want them to do (just with different names):

class ContextTester():
    def __enter__(self):
        self.context = create_context()
        self.context.prepare_context()
        return self.context

    def __exit(self, exc_type, exc_value, exc_traceback):
        self.context.close()

def Test(TestCase):
    def test_action1(self):
        with ContextTester() as context:
            context.action1()
            self.assertTrue(context.check1())

Upvotes: 2

Related Questions