如何通过pytest monkeypatch装饰Stub而非完全替换原函数？

我有如下示例代码：

import pytest

class Helper:
    def function1(self):
        print("helper function called")

class SUT:
    def real_function(self):
        helper = Helper()
        helper.function1()
        print("real function called")

def test_Test(monkeypatch):
    test = SUT()
    def stub_function1(self):
        print("helper is stubbed")

    monkeypatch.setattr(Helper, "function1", stub_function1)
    test.real_function()

运行pytest .\test_scratch_1.py --capture=tee-sys后，按预期输出helper is stubbed和real function called。

但我希望先输出helper function called，再输出helper is stubbed——也就是不完全替换原函数，而是给函数添加后续执行的装饰逻辑。我尝试了以下写法：

import pytest

class Helper:
    def function1(self):
        print("helper function called")

class SUT:
    def real_function(self):
        print("real function called")

def test_Test(monkeypatch):
    test = SUT()
    helper_original = Helper()
    def stub_function1(helper_original_not_stubbed):
        helper_original_not_stubbed.function1()
        print("helper is stubbed")

    monkeypatch.setattr(Helper, "function1", stub_function1(helper_original))
    test.real_function()

但结果和之前一致，请问该如何实现我的需求？

你之前的写法错在：stub_function1(helper_original)是直接调用了这个函数，把调用结果（而非函数本身）赋值给了Helper.function1，自然达不到预期效果。正确的做法是先保存原类方法的引用，再定义一个包装函数，在包装函数里先执行原逻辑，再新增自定义逻辑，最后用monkeypatch替换原方法。

修改后的测试代码如下：

import pytest

class Helper:
    def function1(self):
        print("helper function called")

class SUT:
    def real_function(self):
        helper = Helper()
        helper.function1()
        print("real function called")

def test_Test(monkeypatch):
    test = SUT()
    # 保存原类方法的引用
    original_func = Helper.function1

    # 定义包装函数，保留原方法的self参数，先执行原逻辑再执行新增逻辑
    def wrapped_func(self):
        original_func(self)
        print("helper is stubbed")

    # 替换Helper类的function1为包装函数
    monkeypatch.setattr(Helper, "function1", wrapped_func)
    test.real_function()

运行后会得到预期输出：

helper function called
helper is stubbed
real function called

核心要点 #

• 必须保存原类方法的引用，而非实例方法，因为SUT中会新建Helper实例调用方法
• 包装函数要保留原方法的参数（这里是self），保证调用时的上下文正确
• 给monkeypatch赋值时，要传入包装函数本身，不能传入函数调用后的结果

内容的提问来源于stack exchange，提问作者aldo