如何在Python中Mock带参数的函数并根据参数返回对应值？

你可以通过两种方式实现Variable.get()根据传入参数返回对应值的Mock：

方法一：利用字典作为side_effect #

unittest.mock.Mock的side_effect支持传入字典，此时Mock会把调用时的第一个参数当作字典的key，返回对应的value。示例如下：

import unittest
from unittest.mock import patch, Mock
from your_module import create_file  # 替换成你的代码所在模块

class TestCreateFile(unittest.TestCase):
    def test_create_file_with_mocked_variable(self):
        # 定义参数与返回值的映射
        variable_map = {
            "DIR_NAME": "/tmp/test_dir",
            # 可添加更多参数对应的返回值
        }
        
        # Mock af.data.Variable类的get方法
        with patch("your_module.Variable") as mock_variable:
            mock_variable.get.side_effect = variable_map
            
            # 执行被测函数
            create_file()
            
            # 验证调用是否正确
            mock_variable.get.assert_called_once_with("DIR_NAME")
            # 验证main_client.create是否使用了正确的目录路径
            from your_module import file
            file.Client.create.assert_called_once_with("/tmp/test_dir")

方法二：自定义函数作为side_effect #

如果需要更复杂的逻辑（比如参数匹配规则、动态生成返回值），可以自定义一个函数作为side_effect：

import unittest
from unittest.mock import patch
from your_module import create_file

class TestCreateFile(unittest.TestCase):
    def test_create_file_with_custom_side_effect(self):
        def mock_variable_get(key):
            # 根据传入的key返回对应值
            if key == "DIR_NAME":
                return "/tmp/test_dir"
            elif key == "OTHER_VAR":
                return "other_value"
            # 处理未定义的参数，可返回默认值或抛出异常
            raise ValueError(f"Unknown variable key: {key}")
        
        with patch("your_module.Variable.get", side_effect=mock_variable_get):
            create_file()
            
            # 验证调用
            from your_module import Variable
            Variable.get.assert_called_once_with("DIR_NAME")

注意事项 #

• 确保patch的路径正确：patch("your_module.Variable")中的路径是被测代码中导入Variable的路径，而非原模块af.data.Variable的路径。
• 如果需要验证file.Client的调用，同样需要Mock这个类，比如在测试前添加patch("your_module.file.Client")。

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