是否可以限制模拟函数调用计数？

Question

我在编写单元测试时遇到了问题。这是来自单元测试文件的一个块：

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

我想在单元测试中表明，如果系统中的某些对象具有相似的 short_url，则函数“generate_short_url”不会返回重复值。为此，我使用预定义的返回结果模拟了“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()