3、单元测试基础：使用pytest/unittest对数据转换函数进行测试，Mock外部依赖

单元测试，说白了就是给代码里的最小单元做体检。在数据流水线里，这个最小单元通常就是一个个数据转换函数。你想想看，如果连这些基础函数都不可靠，那整个流水线跑出来的结果你敢信吗？

我个人习惯，在写任何数据转换逻辑之前，先把测试用例写好。这不是什么高深的理论，纯粹是被坑出来的经验。有一次我负责一个实时ETL任务，上游改了字段类型没通知，结果下游全崩了。从那以后，我就养成了「先写测试，再写代码」的习惯。

3.1 为什么数据转换函数需要单元测试

数据转换函数，比如清洗空值、格式转换、字段映射、聚合计算，它们的特点是：

• 输入输出明确：给什么数据，返回什么结果，逻辑清晰
• 容易被修改：业务规则一变，这些函数就得跟着改
• 复用频率高：同一个函数可能在多个流水线里被调用

没有单元测试，你改了一个函数，根本不知道会不会影响其他地方。有了测试，你改完代码跑一遍，绿了就是没问题，红了就知道哪里出事了。

3.2 pytest vs unittest：我选哪个？

Python 里有两个主流的单元测试框架：unittest 是标准库自带的，pytest 是第三方但更流行的。我个人的建议是：

对比维度	unittest	pytest
上手难度	需要写类、继承 TestCase	直接写函数，更简洁
断言方式	self.assertEqual() 等	直接用 assert 语句
Fixture 管理	setUp / tearDown	@pytest.fixture，更灵活
参数化测试	需要自己写循环或子类	@pytest.mark.parametrize
Mock 支持	unittest.mock	pytest-mock 插件
报告输出	基本够用	更美观，支持插件扩展

我一般用 pytest 写新项目，因为它写起来快，读起来也清楚。但如果团队里有人坚持用 unittest，也没问题——两者可以混用，pytest 能直接运行 unittest 写的测试用例。

3.3 实战：测试一个数据清洗函数

假设我们有一个清洗函数，用来处理用户输入的手机号：

# data_cleaner.py
import re

def clean_phone_number(phone: str) -> str:
    """
    清洗手机号：去空格、去横线、去国家码+86，只保留11位数字
    """
    if not phone:
        raise ValueError("手机号不能为空")
    
    # 去空格和横线
    phone = phone.strip().replace(" ", "").replace("-", "")
    
    # 去掉 +86 或 86 前缀
    if phone.startswith("+86"):
        phone = phone[3:]
    elif phone.startswith("86") and len(phone) > 11:
        phone = phone[2:]
    
    # 只保留数字
    phone = re.sub(r"\D", "", phone)
    
    if len(phone) != 11:
        raise ValueError(f"手机号长度不正确: {phone}")
    
    return phone

用 pytest 写测试：

# test_data_cleaner.py
import pytest
from data_cleaner import clean_phone_number

def test_clean_normal_phone():
    """正常手机号，带空格和横线"""
    result = clean_phone_number("138 1234 5678")
    assert result == "13812345678"

def test_clean_with_country_code():
    """带 +86 前缀"""
    result = clean_phone_number("+86 138-1234-5678")
    assert result == "13812345678"

def test_clean_empty_input():
    """空输入应该抛异常"""
    with pytest.raises(ValueError, match="不能为空"):
        clean_phone_number("")

def test_clean_invalid_length():
    """长度不对应该抛异常"""
    with pytest.raises(ValueError, match="长度不正确"):
        clean_phone_number("12345")

跑一下测试：

$ pytest test_data_cleaner.py -v
============================= test session starts ==============================
collected 4 items

test_data_cleaner.py::test_clean_normal_phone PASSED
test_data_cleaner.py::test_clean_with_country_code PASSED
test_data_cleaner.py::test_clean_empty_input PASSED
test_data_cleaner.py::test_clean_invalid_length PASSED

============================== 4 passed in 0.03s ===============================

嗯，全绿。但别高兴太早，这只是最基础的场景。实际项目中，边界情况多得很。

3.4 Mock 外部依赖：让测试跑得更快、更稳

数据转换函数经常要调用外部服务，比如查数据库、调 API、读文件。如果测试里真的去连数据库，那问题就来了：

• 测试跑得慢（网络延迟）
• 测试不稳定（数据库挂了就全红）
• 测试有副作用（真的写了数据进去）

解决办法就是 Mock——用假的对象替换真实的外部依赖。

举个例子，我们有个函数需要根据用户 ID 从数据库查手机号，然后清洗：

# user_service.py
from data_cleaner import clean_phone_number

def get_clean_phone(user_id: int, db_connection) -> str:
    """从数据库获取用户手机号并清洗"""
    cursor = db_connection.cursor()
    cursor.execute("SELECT phone FROM users WHERE id = ?", (user_id,))
    row = cursor.fetchone()
    if not row:
        raise ValueError(f"用户 {user_id} 不存在")
    return clean_phone_number(row[0])

用 unittest.mock 来测试：

# test_user_service.py
import pytest
from unittest.mock import MagicMock
from user_service import get_clean_phone

def test_get_clean_phone_success():
    """模拟数据库返回正常数据"""
    # 创建一个模拟的数据库连接
    mock_db = MagicMock()
    mock_cursor = mock_db.cursor.return_value
    mock_cursor.fetchone.return_value = ("138 1234 5678",)
    
    result = get_clean_phone(1, mock_db)
    assert result == "13812345678"
    
    # 验证 SQL 语句是否正确执行
    mock_cursor.execute.assert_called_once_with(
        "SELECT phone FROM users WHERE id = ?", (1,)
    )

def test_get_clean_phone_user_not_found():
    """模拟用户不存在"""
    mock_db = MagicMock()
    mock_cursor = mock_db.cursor.return_value
    mock_cursor.fetchone.return_value = None
    
    with pytest.raises(ValueError, match="不存在"):
        get_clean_phone(999, mock_db)

你看，用了 Mock 之后，测试不需要真的连数据库，跑起来飞快，而且结果完全可控。

3.5 用 pytest-mock 让 Mock 更简洁

pytest 有个插件叫 pytest-mock，用起来比 unittest.mock 更顺手：

# 安装：pip install pytest-mock

def test_get_clean_phone_with_mock(mocker):
    """使用 pytest-mock 的 mocker fixture"""
    mock_db = mocker.MagicMock()
    mock_cursor = mock_db.cursor.return_value
    mock_cursor.fetchone.return_value = ("139 8765 4321",)
    
    result = get_clean_phone(2, mock_db)
    assert result == "13987654321"

mocker 会自动清理 Mock，不会影响其他测试用例。我个人觉得这比在每个测试函数里手动创建 MagicMock 要干净得多。

3.6 本章知识体系

下面这张图总结了单元测试在数据流水线中的核心逻辑：

3.7 避坑指南

最后，分享几个我踩过的坑：

• 不要 Mock 自己写的函数：如果你 Mock 了正在测试的函数本身，那测试就失去了意义。Mock 只针对外部依赖。
• 注意 Mock 的副作用：有些 Mock 对象会记住调用记录，如果测试用例之间不隔离，可能会互相影响。用 pytest 的 fixture 或 mocker 可以避免这个问题。
• 测试覆盖率不是万能的：我曾经追求 100% 覆盖率，结果发现很多测试只是「为了覆盖而覆盖」，根本没测到关键逻辑。重点应该放在核心路径和边界条件上。
• 测试代码也是代码：别把测试代码写得乱七八糟。好的测试代码应该像文档一样，读一遍就知道被测函数是干什么的。