python - 在单个模拟变量中跨不同文件模拟相同的 Python 函数

Question

假设我有以下 python 文件:

# source.py
def get_one():
    return 1

# file1.py
from source import get_one
def func1():
    return get_one()

# file2.py
from source import get_one
def func2():
    return get_one()

# script.py
from file1 import func1
from file2 import func2
def main(a, b):
    count = 0
    for _ in range(a):
        count += func1()
    for _ in range(b):
        count += func2()
    return count

我知道我可以使用以下设置在 main.py 中模拟 get_one():

def test_mock():
    with (
        patch("file1.get_one") as mock1,
        patch("file2.get_one") as mock2,
    ):
        main(2, 3)
    assert mock1.call_count + mock2.call_count == 5

但是，如果需要在许多文件中模拟 get_one() ，这会变得越来越冗长且难以阅读。我希望能够在一个模拟变量中模拟出它的所有位置。像这样的东西:

# this test fails, I'm just showing what this ideally would look like
def test_mock():
    with patch("file1.get_one", "file2.get_one") as mock:
        main(2, 3)
    assert mock.call_count == 5

有办法做到这一点还是我需要使用多个模拟？

注意，我知道我无法模拟函数的定义位置，例如补丁(“source.get_one”)。

Answer 1

patch 接受 new 作为用于修补目标的对象:

def test_mock():
    with (
        patch("file1.get_one") as mock,
        patch("file2.get_one", new=mock),
    ):
        main(2, 3)
    assert mock.call_count == 5, mock.call_count

您可以编写一个辅助上下文管理器:

import contextlib
from unittest.mock import DEFAULT, patch


@contextlib.contextmanager
def patch_same(target, *targets, new=DEFAULT):
    with patch(target, new=new) as mock:
        if targets:
            with patch_same(*targets, new=mock):
                yield mock
        else:
            yield mock

用法:

def test_mock():
    with patch_same("file1.get_one", "file2.get_one") as mock:
        main(2, 3)
    assert mock.call_count == 5