unittest mock replace/reset mocked function when patching an object

Question

I need to use unittest.mock.patch.object to mock an external method that may fail sometimes. In the test, the method shall raise some errors and then return to the original behaviour. Note that the behaviour I want to reproduce is by far more complicated than return 'bar' so I cannot just copy the code in Bar.some_method_that_may_fail:

import unittest
from unittest.mock import patch


class Bar(object):

    def some_method_that_may_fail(self):
        return "bar"


class Foo(object):
    bar = None

    def retry_method(self):
        try:
            self.__class__.bar = Bar().some_method_that_may_fail()
        except KeyError:
            self.retry_method()


class TestRetry(unittest.TestCase):

    def setUp(self):
        self.instance = Foo()

    def test_retry(self):
        # raise KeyError the first 5 calls to the function and then work normally
        errors_list = [KeyError("")] * 5

        def raise_errors(*_):
            if errors_list:
                errors_list.pop(0)
            # TODO: return to original behaviour

        with patch.object(Bar, 'some_method_that_may_fail', new=raise_errors) as mocked:
            self.instance.retry_method()
        self.assertEqual(self.instance.bar, 'bar')


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

As far as I understand, you want to test `retry_method`, but you can't test it if you are mocking it. Also, this method does not raise an exception, it _catches_ it. What you should mock instead is the stuff in the `try` part (e.g. whatever `self.bar = "bar"` stands for), and _that_ should raise the exception. — MrBean Bremen, Oct 28 '21 at 17:36
@MrBeanBremen you're totally right, I corrected the test because it was wrong. — ronkov, Oct 29 '21 at 12:12

score 3 · Accepted Answer · answered Oct 29 '21 at 17:52

To return different values on subsequent invocations, you can use side_effect. Passing an array of values and/or exceptions will return these values/raise these exceptions in subsequent calls, in your case the exception instances and the result of an original call (if that is what you need). So your test could look something like this:

class TestRetry(unittest.TestCase):

    def setUp(self):
        self.instance = Foo()

    def test_retry(self):
        original = Bar.some_method_that_may_fail  # save the original
        with patch(__name__ + '.Bar') as mocked:
            bar = mocked.return_value
            side_effect = ([KeyError()] * 5) + [original(bar)]
            bar.some_method_that_may_fail.side_effect = side_effect
            self.instance.retry_method()
            self.assertEqual(6, mocked.call_count)
        self.assertEqual('bar', self.instance.bar)

A few notes:

I replaced mock.patch.object with mock.patch, because you don't have an object to patch (Bar is instantiated inside the tested function, and you need to patch the instance, not the class)
using __name__ + '.Bar' for patching is because the tested function is in the same module as the test - in the real code this has to be replaced with the correct module path
I added a check for call_count to ensure that the method has indeed been called 6 times

Another thing: you have an error in your Foo, probably because you dumbed it down for the example. foo is a class variable, but you are setting an instance variable of the same name. You need instead:

class Foo:
    bar = None

    def retry_method(self):
        try:
            self.__class__.bar = Bar().some_method_that_may_fail()
        except KeyError:
            self.retry_method()

(note the self.__class__.bar)

unittest mock replace/reset mocked function when patching an object

1 Answers1