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

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

1 个答案: