如何用unittest仅修改my_func输入'a'的返回值，其余保持原样？

要实现这个部分mock的需求，我们可以借助unittest.mock.patch的side_effect参数，只修改输入为'a'时的返回值，其余情况保留原函数的逻辑。

具体代码实现 #

先导入必要的模块和工具：

import unittest
from unittest.mock import patch
from apps.my_module import my_func

然后编写测试类和测试方法：

class TestMyFunc(unittest.TestCase):
    def test_partial_mock_my_func(self):
        # 定义自定义函数，控制mock的行为逻辑
        def custom_handler(name):
            if name == 'a':
                return 'name is AA'
            # 其他情况直接调用原函数
            return my_func(name)
        
        # 使用patch替换目标函数，将side_effect设为自定义函数
        with patch('apps.my_module.my_func', side_effect=custom_handler):
            # 验证各场景的返回值是否符合预期
            self.assertEqual(my_func('a'), 'name is AA')
            self.assertEqual(my_func('b'), 'name is b')
            self.assertEqual(my_func('c'), 'not defined')

if __name__ == '__main__':
    unittest.main()

原理说明 #

side_effect允许我们传入一个自定义函数，当被mock的函数被调用时，会优先执行这个自定义函数并返回它的结果。我们在custom_handler里只针对name='a'的场景修改返回值，其他场景直接调用原函数，刚好满足“只改特定输入，其余保留原样”的需求。

另一种实现方式（借助wraps） #

如果想更清晰地保留原函数的引用，也可以用wraps参数结合side_effect：

class TestMyFunc(unittest.TestCase):
    def test_partial_mock_with_wraps(self):
        with patch('apps.my_module.my_func', wraps=my_func) as mock_func:
            def custom_handler(name):
                if name == 'a':
                    return 'name is AA'
                # 调用被wraps包裹的原函数
                return mock_func.__wrapped__(name)
            
            mock_func.side_effect = custom_handler
            
            # 同样的验证逻辑
            self.assertEqual(my_func('a'), 'name is AA')
            self.assertEqual(my_func('b'), 'name is b')
            self.assertEqual(my_func('c'), 'not defined')

两种方式都能达到目的，第一种更直接简洁，第二种在需要处理更复杂的mock逻辑时会更灵活。

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