如何编写Python unit mock测试校验对比列表的get_new_sources方法

核心思路 #

get_new_sources的逻辑仅依赖get_sources和get_pds_id两个方法的返回值，不需要对接真实数据库，直接Mock这两个方法的返回值，构造不同测试场景验证逻辑正确性即可，不需要使用真实业务数据。

修正原代码笔误 #

你提供的get_new_sources方法存在未定义变量的笔误，先修正后再进行测试：

class CreateSources():
    def __init__(self):
        pass

    def get_sources(self):
        return list_item

    def get_pds_id(self):
        return list_item

    def get_new_sources(self):
        sources_db_list = self.get_sources()
        # 原代码误用未定义的dbt_sources_db_list，替换为上面赋值的sources_db_list
        sources_db_list = list(dict.fromkeys(sources_db_list))
        db_list = self.get_pds_id()
        # 原代码误用未定义的dremio_db_path，替换为上面赋值的db_list
        db_list = list(dict.fromkeys(db_list))
        # 列表推导式里的变量也同步修正
        new_sources = [a for a in sources_db_list if a not in [x.lower() for x in db_list]]
        return new_sources

完整测试代码实现 #

我们用unittest.mock模拟两个依赖方法的返回值，覆盖三类常见测试场景：

import unittest
from unittest import mock
# 导入你自己的类，如果类名是DremioCreateUpdateSources自行替换
from create_sources import CreateSources

class TestCreateSources(unittest.TestCase):
    # 按照方法路径mock两个依赖的方法，注意mock的参数顺序从下到上
    @mock.patch.object(CreateSources, 'get_pds_id')
    @mock.patch.object(CreateSources, 'get_sources')
    def test_has_new_sources(self, mock_get_sources, mock_get_pds_id):
        # 场景1：存在新增源
        mock_get_sources.return_value = ["source1", "source2", "source3"]
        mock_get_pds_id.return_value = ["SOURCE1", "SOURCE2"]
        cs = CreateSources()
        self.assertEqual(cs.get_new_sources(), ["source3"])

    @mock.patch.object(CreateSources, 'get_pds_id')
    @mock.patch.object(CreateSources, 'get_sources')
    def test_no_new_sources(self, mock_get_sources, mock_get_pds_id):
        # 场景2：无新增源
        mock_get_sources.return_value = ["source1", "source2"]
        mock_get_pds_id.return_value = ["SOURCE1", "SOURCE2"]
        cs = CreateSources()
        self.assertEqual(cs.get_new_sources(), [])

    @mock.patch.object(CreateSources, 'get_pds_id')
    @mock.patch.object(CreateSources, 'get_sources')
    def test_duplicate_source_in_input(self, mock_get_sources, mock_get_pds_id):
        # 场景3：输入列表有重复值，验证去重逻辑正常
        mock_get_sources.return_value = ["source1", "source1", "source2"]
        mock_get_pds_id.return_value = ["SOURCE1", "SOURCE1"]
        cs = CreateSources()
        self.assertEqual(cs.get_new_sources(), ["source2"])

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

逻辑说明 #

• 用@mock.patch.object装饰器模拟类的内部方法，避免触发真实的数据库查询逻辑
• 每个测试用例独立构造模拟返回值，覆盖正常、无新增、输入重复等边界场景
• 直接通过断言判断返回结果是否符合预期，无需手动print校验测试结果

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