41. 单元测试

写完功能代码只是第一步。你怎么确定它在各种情况下都能正常工作？手动在交互式环境里试几次？输入变多后根本试不过来。

单元测试（unit testing）就是自动化回答这个问题的手段。它把代码拆成最小单元（通常是函数或方法），为每个单元编写自动化测试。运行测试套件，几分钟就能验证整个项目的正确性。

unittest 入门

Python 标准库自带 unittest 模块，用法跟 Java 的 JUnit 类似：

import unittest

def add(a, b):
    return a + b

class TestMath(unittest.TestCase):
    def test_add_integers(self):
        self.assertEqual(add(1, 2), 3)

    def test_add_floats(self):
        self.assertAlmostEqual(add(0.1, 0.2), 0.3, places=7)

    def test_add_negative(self):
        self.assertEqual(add(-1, 1), 0)

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

运行这个文件，如果所有测试通过，会看到几个点：

...
----------------------------------------------------------------------
Ran 3 tests in 0.001s

OK

常用断言方法

unittest.TestCase 提供了一系列断言方法：

方法	用途
assertEqual(a, b)	a == b
assertNotEqual(a, b)	a != b
assertTrue(x)	bool(x) is True
assertFalse(x)	bool(x) is False
assertIsNone(x)	x is None
assertIn(a, b)	a in b
assertIsInstance(a, b)	isinstance(a, b)
assertRaises(exc)	上下文管理器，验证抛出了指定异常

assertRaises 的用法稍特殊：

def divide(a, b):
    if b == 0:
        raise ValueError("不能除以零")
    return a / b

class TestDivide(unittest.TestCase):
    def test_divide_by_zero(self):
        with self.assertRaises(ValueError):
            divide(10, 0)

setUp 与 tearDown

很多测试需要准备环境（比如创建临时文件、连接数据库）和清理环境。setUp 在每个测试方法执行前运行，tearDown 在每个测试方法执行后运行：

import unittest
import tempfile
import os

class TestFileOps(unittest.TestCase):
    def setUp(self):
        self.tmpdir = tempfile.mkdtemp()
        self.testfile = os.path.join(self.tmpdir, 'test.txt')
        with open(self.testfile, 'w') as f:
            f.write('hello')

    def tearDown(self):
        import shutil
        shutil.rmtree(self.tmpdir)

    def test_read(self):
        with open(self.testfile, 'r') as f:
            content = f.read()
        self.assertEqual(content, 'hello')

    def test_exists(self):
        self.assertTrue(os.path.exists(self.testfile))

两个测试方法共享了同一份准备和清理逻辑，互不干扰。

测试类级别的准备：setUpClass / tearDownClass

如果一组测试共享昂贵的资源（比如数据库连接），不想每个方法都重建一次，可以用类方法：

class TestDatabase(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.conn = create_db_connection()

    @classmethod
    def tearDownClass(cls):
        cls.conn.close()

    def test_query(self):
        result = self.conn.execute("SELECT 1")
        self.assertEqual(result.fetchone()[0], 1)

setUpClass 在所有测试方法开始前执行一次，tearDownClass 在所有方法结束后执行一次。注意这两个方法要加 @classmethod 装饰器。

Mock：模拟外部依赖

单元测试的原则是「测一个单元，隔离其他单元」。如果被测函数调用了数据库、网络 API 或文件系统，测试会变得又慢又脆弱。

unittest.mock 模块可以创建「假装的对象」，替代真实的外部依赖：

from unittest.mock import Mock, patch
import unittest

def get_weather(city):
    import requests
    resp = requests.get(f"https://api.weather.com/{city}")
    return resp.json()['temperature']

class TestWeather(unittest.TestCase):
    @patch('requests.get')
    def test_get_weather(self, mock_get):
        mock_get.return_value.json.return_value = {'temperature': 25}

        temp = get_weather('北京')
        self.assertEqual(temp, 25)
        mock_get.assert_called_with('https://api.weather.com/北京')

@patch('requests.get') 把 requests.get 替换成一个 Mock 对象。测试中它不会真的发网络请求，而是返回你预设的值。

运行测试的方式

除了 if __name__ == '__main__': unittest.main()，还可以用命令行：

# 运行单个测试文件
python -m unittest test_math.py

# 运行指定类
python -m unittest test_math.TestMath

# 运行指定方法
python -m unittest test_math.TestMath.test_add_integers

# 详细输出
python -m unittest -v test_math.py

-v 参数会打印每个测试方法的名称和结果，排查问题时更直观。

小结

• 测试类继承 unittest.TestCase，测试方法以 test_ 开头。
• 用 setUp/tearDown 管理每个测试的环境，用 setUpClass/tearDownClass 管理类级别的资源。
• unittest.mock 帮你隔离外部依赖，让测试又快又稳。
• 测试命令用 python -m unittest，-v 看详细信息。

写测试一开始会觉得麻烦，但项目越大，测试带来的安全感越强。改了一个函数，跑一遍测试，五分钟就知道有没有破坏其他地方。这种信心，是手动测试给不了的。

---