Python 单元测试入门:unittest 与 pytest

改了一行代码,三个月前的功能悄悄坏了,而你直到上线才发现——这种事每个程序员都经历过。单元测试就是给代码上的保险:把每个函数的行为固化成自动化检查,改完代码跑一遍,坏没坏立刻知道。本文用标准库 unittest 讲透测试的三大件(断言、夹具、组织),最后对比更简洁的 pytest。

1. 什么是单元测试

单元测试针对的是代码里最小的逻辑单元——通常是一个函数。做法很朴素:给定输入,断言输出符合预期。假设有个计算折扣价的函数:

def discounted_price(price: float, discount: float = 0.9) -> float:
    """返回打折后的价格,discount 是折扣率,如 0.9 表示九折。"""
    if not 0 < discount <= 1:
        raise ValueError("折扣率必须在 (0, 1] 之间")
    return round(price * discount, 2)

这个函数能不能用,靠人肉测几次也行,但下次有人改了它(比如忘了 round),你不会知道。单元测试就是把这些"人肉验证"固化成脚本:每次改动后跑一遍,几百个断言替你盯着。测试不是给代码找麻烦,是给未来的自己省麻烦。

那是不是所有代码都要写测试?行业里有个"测试金字塔"的说法:底层是大量单元测试(测函数)、中间是集成测试(测模块协作)、顶层是少量端到端测试(测完整流程)。对初学者来说,优先级很简单:先给"有明确输入输出、逻辑容易出错"的函数写测试——计算、解析、校验、排序这类纯函数性价比最高;而界面点击、网络请求这类难以稳定的代码,可以等基础测试跑顺了再碰。

2. 第一个测试用例

unittest 的玩法是:写一个继承 unittest.TestCase 的类,类里以 test_ 开头的每个方法就是一个测试用例:

import unittest
from math_utils import discounted_price

class TestDiscountedPrice(unittest.TestCase):
    def test_normal_discount(self):
        self.assertEqual(discounted_price(100), 90.0)

    def test_custom_discount(self):
        self.assertEqual(discounted_price(100, 0.5), 50.0)

    def test_invalid_discount(self):
        with self.assertRaises(ValueError):
            discounted_price(100, 1.5)

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

把上面的 discounted_price 存成 math_utils.py,测试存成 test_math_utils.py,同目录运行 python test_math_utils.py,会看到 OK 和三个用例通过的提示。两个关键规则:测试方法必须以 test_ 开头,否则不会被发现;每个测试之间相互独立,一个失败不影响其他用例继续跑。

运行输出值得学会读:全部通过时打印 Ran 3 tests ... OK;有失败时是 FAIL 并附上断言对比;代码本身抛异常则是 ERROR 并给出完整堆栈。区别很重要——FAIL 是"结果不符合预期",ERROR 是"测试代码自己崩了",后者往往意味着被测函数抛出了没预料到的异常。看到 F、E、. 组成的进度条,别慌,. 是通过,字母才是问题。

3. 断言方法:测试的语言

断言是测试的核心词汇表。最常见的几个:

断言	检查内容
`assertEqual(a, b)`	a 等于 b
`assertTrue(x)` / `assertFalse(x)`	x 为真 / 为假
`assertIn(item, container)`	item 在容器里
`assertAlmostEqual(a, b)`	浮点数近似相等(处理精度)
`assertRaises(Exc)`	抛出指定异常

其中 assertAlmostEqual 是个经典陷阱:浮点数运算有精度误差,0.1 + 0.2 == 0.3 是 False,直接 assertEqual 必挂,必须用近似断言。看一个综合例子:

另外注意 assertTrue 的用法边界:它只检查"真值",所以 assertTrue(x == 1) 不如直接 assertEqual(x, 1)——后者失败时能告诉你 x 实际是多少,前者只告诉你"不成立"。同样的道理,assertTrue(len(items) > 0) 可以换成更语义化的 assertGreater(len(items), 0) 或 assertTrue(items)。断言方法选得越精准,失败时得到的诊断信息越多,排查越快。

import unittest

class TestTips(unittest.TestCase):
    def test_float(self):
        # assertEqual(0.1 + 0.2, 0.3)  # 会失败!
        self.assertAlmostEqual(0.1 + 0.2, 0.3)  # 通过

    def test_container(self):
        self.assertIn("py", "python")
        self.assertNotIn("x", "python")

    def test_raises(self):
        with self.assertRaises(ValueError):
            int("不是数字")

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

被注释掉的那行是"错误示范":你可以取消注释跑一次,亲眼看看浮点精度是怎么坑人的。assertRaises 配合 with 语句,专门验证"应该报错的情况",比如参数校验、文件不存在等异常路径。

4. 夹具:setUp 与 tearDown

很多测试用例需要先准备数据。与其在每个方法里重复写,不如用 setUp 统一准备——它在每个测试方法执行前自动调用,tearDown 则在每个测试后清理:

import unittest

class TestScoreBoard(unittest.TestCase):
    def setUp(self):
        # 每个测试开始前都会执行,保证数据是全新的
        self.scores = {"alice": 90, "bob": 85}

    def tearDown(self):
        # 每个测试结束后执行,可在这里释放资源、删临时文件
        self.scores.clear()

    def test_get_score(self):
        self.assertEqual(self.scores["alice"], 90)

    def test_update_score(self):
        self.scores["bob"] = 88
        self.assertEqual(self.scores["bob"], 88)
        # 即便这里改坏了数据,下一个测试的 setUp 也会重建干净数据

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

这套机制保证测试之间互不污染:每个用例拿到的都是 setUp 刚建好的全新状态。注意 setUp 是按"每个测试"执行,不是整个类只执行一次——如果你想让某些重量级准备只做一次(比如连接数据库、加载大文件),用 setUpClass 类方法。

举个现实场景:测试要读写临时文件,setUp 里用 tempfile 建一个,tearDown 里删掉,这样测试跑完磁盘干干净净;测试要连数据库,setUp 里建表、tearDown 里清空数据,保证每次跑都是同一起跑线。记住一条原则:测试必须可重复——同一份代码跑十次,结果应该一模一样。如果测试依赖网络、当前时间、随机数,它就会"时好时坏",这种测试比没有还糟,因为它会消耗你的信任。

5. 组织测试与命令行运行

项目变大后,测试文件要和源码分开管理。推荐结构是把测试放进 tests/ 目录,每个被测模块对应一个测试文件,再用 discover 一键发现并运行所有测试:

project/
├── math_utils.py          # 被测源码
├── main.py
└── tests/
    ├── __init__.py        # 让 tests 成为包,必须存在
    ├── test_math_utils.py
    └── test_main.py

这个结构的核心是"源码与测试分离":被测代码在项目根目录,测试全部收进 tests/,互不干扰。文件命名约定是 test_模块名.py,一眼就知道某个测试在测谁。接下来用命令行批量运行:

# 递归发现 tests/ 下所有测试并运行
python -m unittest discover tests

# 指定运行某个文件,加 -v 显示每个用例的名字
python -m unittest tests.test_math_utils -v

# 运行单个用例
python -m unittest tests.test_math_utils.TestDiscountedPrice.test_normal_discount

discover 会自动查找以 test 开头的文件和以 test_ 开头的方法。-v 参数很实用,能列出每个用例的通过/失败状态,定位问题更快。命令行规则是 python -m unittest 模块路径,用点号代替斜杠,和 import 的写法一致。

那个 tests/__init__.py 空文件容易被人忽略,但非常重要:它让 tests 变成 Python 包,discover 才能用模块路径导入测试文件。如果测试文件要 import 项目根目录的源码,通常还需要把项目根目录加入 PYTHONPATH(比如在项目根目录运行测试、或用 python -m unittest 而不是 python tests/test_x.py)。这些小坑都是入门时最容易卡住的地方,遇到了别怀疑人生,基本都是导入路径问题。

6. pytest:更简洁的写法

unittest 功能齐全但啰嗦:必须写类、必须记一堆 assertXxx 方法名。pytest 把这些简化到极致——普通函数加原生 assert 就行。它需要安装:

pip install pytest

装好后,把测试代码写进 test_discount.py,和 unittest 版本对比着看差异:

# test_discount.py —— 用 pytest 运行:pytest test_discount.py -v
from math_utils import discounted_price

def test_normal_discount():
    assert discounted_price(100) == 90.0

def test_custom_discount():
    assert discounted_price(100, 0.5) == 50.0

def test_invalid_discount():
    import pytest
    with pytest.raises(ValueError):
        discounted_price(100, 1.5)

同样三个测试,代码量少了一半:不用继承 TestCase,断言直接用 Python 的 assert,失败信息却比 unittest 更详细(pytest 会自动展示两边的值)。pytest 的 fixture 机制(用 @pytest.fixture 装饰器)也比 setUp 更灵活,支持按需注入。入门建议:先吃透 unittest 的概念,再切 pytest 提升效率,两个都认识才不会被项目里的旧测试吓到。

pytest 还有两个高频利器值得提前知道:参数化 @pytest.mark.parametrize("输入, 期望", [(1, 2), (2, 4)]) 能把"同一逻辑、多组数据"的用例压成一行,再也不用复制粘贴测试方法;fixture 则支持作用域控制(scope="module" 表示整个模块只准备一次)和依赖注入,比 setUp 家族灵活得多。等你用 unittest 写过几十个用例,再回头看 pytest,会明显感受到"少写样板代码"的快乐。

7. 总结与练习

回顾核心:测试用例 = 准备输入 + 断言输出;unittest 用 TestCase 类组织,test_ 前缀发现用例,setUp/tearDown 管理夹具,discover 批量运行;pytest 用纯函数 + assert 提供更简洁的体验。

练习 1:写一个字符串处理函数(比如去掉首尾空格并转小写),为它补 5 个测试用例,覆盖正常、空串、全大写等边界。
练习 2:用 setUp 准备一份用户字典,写两个互相独立的测试,验证数据不被互相污染。
练习 3:把同一组测试分别用 unittest 和 pytest 实现,对比两种写法的代码量,体会差异。

💡 写测试的黄金顺序:先写测试,再写实现——这叫测试驱动开发(TDD)。哪怕不严格照做,也请记住:新功能上线前,先问自己"如果这个函数坏了,测试能不能第一时间发现"。