如何使用mock.patch修补含yield的生成器类方法？常规方法失效求解

问题场景 #

要修补PassiveJsonMixin类里的生成器方法get_changed_diff_patch，原类代码如下：

class PassiveJsonMixin(JsonMixin):
    """
    passive items that do not have id/key
    """

    def __init__(self, *args, **kwargs):
        JsonMixin.__init__(self, *args, **kwargs)
        # self.build_json()
        pass
    def get_changed_diff_patch(self, parent_hidden=False):
        print ("From PassiveJsonMixin: get_changed_diff_patch")
        return
        yield

试过常规的patch.object方式，但没生效：

patch.object(mymod.PassiveJsonMixin, 'get_changed_diff_patch',wrapper(mymod.PassiveJsonMixin.get_changed_diff_patch) )

可行的修补方案 #

方案1：正确包装生成器方法 #

生成器方法的修补得保证包装后的函数依然返回生成器（或可迭代对象），不能直接换成普通函数。你可以写个正确的包装函数，保留原生成器逻辑，或者加自定义逻辑：

from unittest.mock import patch

def patched_get_changed_diff_patch(original_func):
    def wrapper(self, parent_hidden=False):
        print("Patched: before original generator")
        # 调用原生成器并yield其结果
        yield from original_func(self, parent_hidden)
        print("Patched: after original generator")
        # 可以加自定义的yield内容
        yield {"custom_patch": "added_content"}
    return wrapper

# 用patch.object替换方法
with patch.object(mymod.PassiveJsonMixin, 'get_changed_diff_patch', 
                  patched_get_changed_diff_patch(mymod.PassiveJsonMixin.get_changed_diff_patch)):
    # 测试代码
    obj = mymod.PassiveJsonMixin()
    for item in obj.get_changed_diff_patch():
        print(item)

方案2：用side_effect做简单替换 #

要是需要完全替换原生成器的返回值，用side_effect指定一个新的生成器函数就行：

from unittest.mock import patch

def mock_generator(self, parent_hidden=False):
    print("From mock generator")
    yield {"diff": "mock_diff"}

with patch.object(mymod.PassiveJsonMixin, 'get_changed_diff_patch', side_effect=mock_generator):
    obj = mymod.PassiveJsonMixin()
    for item in obj.get_changed_diff_patch():
        print(item)

方案3：测试用例里用装饰器式patch #

如果是在单元测试中，可以用装饰器形式的patch，同样要保证替换的是生成器函数：

from unittest.mock import patch
import unittest

class TestPassiveJsonMixin(unittest.TestCase):
    @patch.object(mymod.PassiveJsonMixin, 'get_changed_diff_patch')
    def test_patched_generator(self, mock_method):
        # 设置mock返回一个生成器
        mock_method.return_value = (item for item in [{"test": "data"}])
        obj = mymod.PassiveJsonMixin()
        result = list(obj.get_changed_diff_patch())
        self.assertEqual(result, [{"test": "data"}])

原方法不生效的原因 #

你之前的写法问题出在：

• 直接把包装后的函数赋值给被patch的方法，但如果wrapper没正确处理生成器的yield逻辑（比如没返回生成器），就会丢了原方法的生成器特性。
• 另外，patch.object的第三个参数如果是普通函数，会直接替换原方法，但生成器方法得返回可迭代的生成器对象，普通函数没法替代。

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