在單個模擬變量中跨不同文件模擬相同的 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

無論如何要這樣做還是我需要使用多個模擬？

注意，我知道我不能模擬函數的定義位置，例如patch("source.get_one") 。

Answer 1

一種可能的解決方法：在source.py中添加一個間接級別，以便get_one具有可以修補的依賴項：

def get_one():
    return _get_one()

def _get_one():
    return 1

現在您可以在測試中執行此操作：

from script import main
from unittest.mock import patch

def test_mock():
    with patch("source._get_one") as mock1:
        main(2, 3)
    assert mock1.call_count == 5

Answer 2

您可以編寫一個輔助上下文管理器：

import contextlib
from unittest.mock import DEFAULT, patch


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

用法：

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