mardi 13 septembre 2016

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 ?

Aucun commentaire:

Enregistrer un commentaire