如何在Python中Mock yaml.load()和Box()并测试Config类？

问题场景 #

给定如下Config类，其中create_config方法依赖yaml.load()和Box()两个外部库的功能：

import yaml
from box import Box
from yaml import SafeLoader


class Config:
    def set_config_path(self):
        self.path = r"./config/datasets.yaml"
        return self.path

    def create_config(self):
        with open(r"./config/datasets.yaml") as f:
            self.config = Box(yaml.load(f, Loader=SafeLoader))
        return self.config

已编写的unittest测试用例在测试create_config时无法正确验证逻辑，需要Mockyaml.load()和Box()，传入假文件数据后得到预期的假字典配置。原测试用例如下：

import unittest
from unittest.mock import mock_open, patch
from src.utils.config import Config


class TestConfig(unittest.TestCase):
    def setUp(self):
        self.path = r"./config/datasets.yaml"

    def test_set_config_path(self):
        assert Config.set_config_path(self) == self.path

    @patch("builtins.open", new_callable=mock_open, read_data="data")
    def test_create_config(self, mock_file):
        assert open(self.path).read() == "data"

解决方案 #

要Mock外部库的函数/类，核心原则是patch代码实际引用它们的位置（而非库本身的模块），也就是Config类所在模块（src.utils.config）里的yaml.load和Box。具体实现如下：

完整测试用例代码 #

import unittest
from unittest.mock import mock_open, patch, MagicMock
from src.utils.config import Config


class TestConfig(unittest.TestCase):
    def setUp(self):
        self.path = r"./config/datasets.yaml"
        # 预设的假配置数据
        self.mock_yaml_data = {"dataset": "mnist", "batch_size": 32}

    def test_set_config_path(self):
        config = Config()
        assert config.set_config_path() == self.path

    @patch("builtins.open", new_callable=mock_open, read_data="fake_yaml_content")
    @patch("src.utils.config.yaml.load")
    @patch("src.utils.config.Box")
    def test_create_config(self, mock_box, mock_yaml_load, mock_file):
        # 配置yaml.load返回预设的假字典
        mock_yaml_load.return_value = self.mock_yaml_data
        
        # 配置Box：直接返回假字典，或用MagicMock模拟其属性访问行为
        mock_box.return_value = self.mock_yaml_data

        # 执行测试方法
        config = Config()
        result = config.create_config()

        # 验证各Mock对象的调用情况
        mock_file.assert_called_once_with(self.path)
        mock_yaml_load.assert_called_once_with(mock_file.return_value.__enter__(), Loader=Config.create_config.__globals__["SafeLoader"])
        mock_box.assert_called_once_with(self.mock_yaml_data)
        # 验证返回结果符合预期
        assert result == self.mock_yaml_data

关键说明 #

• Patch路径的选择：必须patchsrc.utils.config.yaml.load而非yaml.load，因为Config类是从自身模块的导入中使用这些外部依赖的，这是Python Mock的核心规则——patch where it's used。
• Mock的灵活性：如果需要验证Box的属性访问逻辑（比如config.config.dataset），可以用MagicMock替代直接返回字典，例如mock_box.return_value = MagicMock(**self.mock_yaml_data)。
• 文件上下文的处理：mock_open返回的对象在with语句中会触发__enter__()方法，因此验证yaml.load的参数时，需要用mock_file.return_value.__enter__()模拟打开后的文件对象。

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