Можно ли ограничить количество посмешиваемых вызовов функций?

Question

У меня возникла проблема, когда я пишу модульный тест. Это chunck из блока тестового файла:

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() 

Это кусок кода из файла «prj.apps.app.models.base» (файл, который импортирует функцию «generate_url_string», который дразнят) :

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 

Я хочу показать в тесте на единицу, что функция «generate_short_url» не возвращает повторяющиеся значения, если некоторые объекты в системе имеют схожие short_urls. Я обманул «generate_url_string» с предопределенным результатом возврата для этой цели. Проблема в том, что я не мог ограничить количество вызовов издеваемой функции с этим значением, и в результате код переходит в бесконечный цикл. Я хотел бы назвать свою функцию с заранее заданным результатом ('a1b2c3') только один раз. После этого я хочу, чтобы функция работала как обычно. Что-то вроде этого:

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() 

Но я не вижу каких-либо атрибутов, таких как «» times_to_call в насмешливой библиотеке. Есть ли способ справиться с этим?

Answer 1

Определите генератор, который сначала дает фиксированное значение, а затем возвращает возвращаемое значение реальной функции (которое передается как аргумент, чтобы избежать вызова исправленного значения).

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

Затем используйте генератор в качестве побочного эффекта исправленной функции.

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()