Roman Storozhenko Roman Storozhenko - 2 months ago 7
Python Question

Is it possible to limit mocked function calls count?

I have encountered a problem when I write a unit test. This is a chunck from an unit test file:

main.obj = MainObj.objects.create(short_url="a1b2c3")

with unittest.mock.patch('prj.apps.app.models.base.generate_url_string', return_value="a1b2c3") as mocked_generate_url_string:
obj.generate_short_url()


This is a chunk of code from the file 'prj.apps.app.models.base' (file which imports function 'generate_url_string' which is being mocked):

from ..utils import generate_url_string
.....................
def generate_short_url(self):
short_url = generate_url_string()
while MainObj.objects.filter(short_url=short_url).count():
short_url = generate_url_string()

return short_url


I want to show in the unit test that the function 'generate_short_url' doesn't return repeated values if some objects in the system have similar short_urls. I mocked 'generate_url_string' with predefined return result for this purpose.
The problem is that I couldn't limit number of calls of mocked function with this value, and as a result the code goes to an infinite loop.
I would like to call my function with predefined result ('a1b2c3') only once. After that I want function to work as usual. Something like this:

with unittest.mock.patch('prj.apps.app.models.base.generate_url_string', return_value="a1b2c3", times_to_call=1) as mocked_generate_url_string:
obj.generate_short_url()


But I see no any attributes like 'times_to_call' in a mocking library.
Is there any way to handle that ?

Answer

Define a generator that first yields the fixed value, then yields the return value of the real function (which is passed as an argument to avoid calling the patched value).

def mocked(x):
    yield "a1b2c3"
    while True:
        yield x()

Then, use the generator as the side effect of the patched function.

with unittest.mock.patch(
       'prj.apps.app.models.base.generate_url_string',
       side_effect=mocked(prj.apps.app.models.base.generate_url_string)) as mocked_generate_url_string:
    obj.generate_short_url()