How to override a mock for an individual test within a class that already has a mock

I have a test class that has a mock decorator, and several tests. Each test receives the mock, because mock is defined on the class level. Great. Here's what it looks like:

@mock.patch("foo", bar)
class TestMyThing(TestCase):
  def test_A(self):
    assert something

  def test_B(self):
    assert something

  def test_C(self):
    assert something

  def test_D(self):
    assert something

I now want test_D to get a have a different value mocked for foo. I first try:

@mock.patch("foo", bar)
class TestMyThing(TestCase):
  def test_A(self):
    assert something

  def test_B(self):
    assert something

  def test_C(self):
    assert something

  @mock.patch("foo", baz)
  def test_D(self):
    assert something

This doesn't work. Currently to get unittest to take the mock.patch that decorates test_D, I have to remove the mock.patch that decorates the class. This means creating lots of DRY and doing the following:

class TestMyThing(TestCase):
  @mock.patch("foo", bar)
  def test_A(self):
    assert something

  @mock.patch("foo", bar)
  def test_B(self):
    assert something

  @mock.patch("foo", bar)
  def test_C(self):
    assert something

  @mock.patch("foo", baz)
  def test_D(self):
    assert something

This is non ideal due to DRY boilerplate, which makes it error prone and violates open-closed principle. Is there a better way to achieve the same logic?

Solution

Yes! You can leverage the setUp/tearDown methods of the unittest.TestCase and the fact that unittest.mock.patch in its "pure" form (i.e. not as context manager or decorator) returns a "patcher" object that has start/stop methods to control when exactly it should do its magic.

You can call on the patcher to start inside setUp and to stop inside tearDown and if you keep a reference to it in an attribute of your test case, you can easily stop it manually in selected test methods. Here is a full working example:

from unittest import TestCase
from unittest.mock import patch


class Foo:
    @staticmethod
    def bar() -> int:
        return 1


class TestMyThing(TestCase):
    def setUp(self) -> None:
        self.foo_bar_patcher = patch.object(Foo, "bar", return_value=42)
        self.mock_foo_bar = self.foo_bar_patcher.start()
        super().setUp()

    def tearDown(self) -> None:
        self.foo_bar_patcher.stop()
        super().tearDown()

    def test_a(self):
        self.assertEqual(42, Foo.bar())

    def test_b(self):
        self.assertEqual(42, Foo.bar())

    def test_c(self):
        self.assertEqual(42, Foo.bar())

    def test_d(self):
        self.foo_bar_patcher.stop()
        self.assertEqual(1, Foo.bar())

This patching behavior is the same, regardless of the different variations like patch.object (which I used here), patch.multiple etc.

Note that for this example it is not necessary to keep a reference to the actual MagicMock instance generated by the patcher in an attribute, like I did with mock_foo_bar. I just usually do that because I often want to check something against the mock instance in my test methods.

It is worth mentioning that you can also use setUpClass/tearDownClass for this, but then you need to be careful with re-starting the patch, if you stop it because those methods are (as the name implies) only called once for each test case, whereas setUp/tearDown are called once before/after each test method.

PS: The default implementations of setUp/tearDown on TestCase do nothing, but it is still good practice IMO to make a habit of calling the superclass' method, unless you deliberately want to omit that call.