为 Python 数据处理脚本编写可维护的单元测试

从一次翻车说起

上个月有个跑了小半年的数据清洗脚本突然出了问题：上游把 order_date 这一列的格式从 2024-01-01 换成了 2024/1/1，脚本没报错，就是默默吐出一堆空值，等业务那边找过来的时候，脏数据已经进了库。排查花了一下午，改脚本只花了两分钟，这就很难受了。

痛定思痛，决定给这些脚本补上单元测试。结果一搜教程，清一色都在教你怎么测 add(1, 2) == 3，可我的脚本里全是在读写文件、调 pandas，这玩意儿到底怎么测啊... 讲“测试金字塔”的多，讲人话的少。下面是实测下来比较顺手的一套做法，不追求优雅，能防住翻车就行。

核心就一招：把读写文件和数据处理拆开

大部分数据处理脚本长这样：

def main():
    df = pd.read_csv("data/raw_orders.csv")
    df["amount"] = df["amount"].fillna(0)
    df["order_date"] = pd.to_datetime(df["order_date"])
    df = df[df["status"] != "cancelled"]
    df.to_csv("data/clean_orders.csv", index=False)

这种是没法测的，因为逻辑和 I/O 糊在一起，一调用就得真去读文件。改造很简单，把纯逻辑抽出来，让它只接收 DataFrame、返回 DataFrame：

def clean_orders(df: pd.DataFrame) -> pd.DataFrame:
    df = df.copy()
    df["amount"] = df["amount"].fillna(0)
    df["order_date"] = pd.to_datetime(df["order_date"], errors="coerce")
    return df[df["status"] != "cancelled"]

def main():
    df = pd.read_csv("data/raw_orders.csv")
    clean_orders(df).to_csv("data/clean_orders.csv", index=False)

main 里那几行本来就没啥逻辑，出了问题基本是文件不存在或者路径写错，跑一下就知道，不用测。真正容易埋雷的是 clean_orders 里这些转换规则，测它就够了。注意 df.copy() 别省，不然函数里改的就是传进来的那个 DataFrame，几个测试跑下来会互相污染（我就吃过这个亏）。

别一上来就想着整个项目大重构，先从最容易出错的那个函数开始抽，抽一个测一个。

用 parametrize 把边界情况铺满

数据处理最容易死在边角料上：空字符串、None、带空格带货币符号的数字。比如把 ¥12.5 转成 12.5 的这种小函数，一个用例肯定不够，推荐用 pytest 的参数化，一次把各种脏输入列全：

import pytest
from myscript.clean import parse_price

@pytest.mark.parametrize("raw, expected", [
    ("¥12.5", 12.5),
    (" 3 ", 3.0),
    ("", None),
    (None, None),
])
def test_parse_price(raw, expected):
    assert parse_price(raw) == expected

以后上游又整出新格式，直接往列表里加一行就行。挂了的时候 pytest 会把具体是哪组参数挂了打印在报错里，看一眼就知道。顺便说一句，测试名直接写人话，别偷懒写成 test_1，不然过俩月连你自己都不知道这测试在测啥。

真要测文件读写的话，用 tmp_path

有些函数就是得碰文件，比如最后负责落盘的那步。不用费劲去 mock 整个文件系统，pytest 自带一个 tmp_path，每个测试给你一个临时目录，用完自动清理：

def test_save_csv(tmp_path):
    out = tmp_path / "result.csv"
    save_csv(pd.DataFrame({"a": [1, 2]}), out)
    assert out.exists()
    assert pd.read_csv(out)["a"].tolist() == [1, 2]

另外，断言两个 DataFrame 别直接写 df1 == df2，那玩意儿返回的还是个 DataFrame，永远为真，要用 pd.testing.assert_frame_equal，遇到 dtype 对不上就加个 check_dtype=False，不用太较真。

最后

目录结构不用花哨，项目里建个 tests/，写完跑一句 pytest，绿了再提交。不用一上来就上 CI、追覆盖率，先把上次翻车的那种问题写成一个测试用例，比什么都实在。

不出问题的话，下次上游再偷偷改格式，你两分钟就能知道。