如何使用 unittest 为 Python 函数编写测试用例

更新于 2026-01-11

DavidMuller, Kathryn Hancox 2020-10-01

简介

Python 标准库中包含 unittest 模块，可帮助你为 Python 代码编写和运行测试。

使用 unittest 编写的测试有助于发现程序中的错误，并在你随着时间推移修改代码时防止回归问题。遵循测试驱动开发（TDD）的团队可以利用 unittest 来确保所有编写的代码都有一组对应的测试。

在本教程中，你将使用 Python 的 unittest 模块为一个函数编写测试。

定义 TestCase 子类

unittest 模块提供的最重要类之一是 TestCase。TestCase 为测试函数提供了通用的框架结构。让我们看一个例子：

test_add_fish_to_aquarium.py

import unittest

def add_fish_to_aquarium(fish_list):
    if len(fish_list) > 10:
        raise ValueError("A maximum of 10 fish can be added to the aquarium")
    return {"tank_a": fish_list}


class TestAddFishToAquarium(unittest.TestCase):
    def test_add_fish_to_aquarium_success(self):
        actual = add_fish_to_aquarium(fish_list=["shark", "tuna"])
        expected = {"tank_a": ["shark", "tuna"]}
        self.assertEqual(actual, expected)

首先我们导入 unittest 模块以供代码使用。然后定义要测试的函数——这里是 add_fish_to_aquarium。

该函数接收一个名为 fish_list 的鱼类列表，如果列表元素超过 10 个，则抛出异常；否则返回一个字典，将鱼缸名称 "tank_a" 映射到给定的 fish_list。

接着我们定义了一个名为 TestAddFishToAquarium 的类，它是 unittest.TestCase 的子类。其中定义了一个方法 test_add_fish_to_aquarium_success，它调用 add_fish_to_aquarium 函数并传入特定输入，验证实际返回值是否与预期一致。

现在我们已经定义了一个包含测试的 TestCase 子类，接下来介绍如何执行这个测试。

执行 TestCase

在上一节中，我们创建了名为 TestAddFishToAquarium 的 TestCase 子类。在同一目录下运行以下命令即可执行测试：

python -m unittest test_add_fish_to_aquarium.py

我们通过 python -m unittest 调用了 Python 的 unittest 模块，并将包含 TestAddFishToAquarium 测试用例的文件路径作为参数传入。

执行后会看到类似如下输出：

.
----------------------------------------------------------------------
Ran 1 test in 0.000s

OK

unittest 模块运行了我们的测试，并告诉我们测试通过。输出第一行的单个 . 表示测试已通过。

注意：TestCase 会将任何以 test 开头的方法识别为测试方法。例如，def test_add_fish_to_aquarium_success(self) 会被识别为测试并执行；而 def example_test(self) 则不会被识别为测试，因为它不以 test 开头。只有以 test 开头的方法才会在运行 python -m unittest ... 时被执行和报告。

尝试一个失败的测试

现在我们修改测试方法中的某一行，人为制造一个失败：

test_add_fish_to_aquarium.py（修改后）

import unittest

def add_fish_to_aquarium(fish_list):
    if len(fish_list) > 10:
        raise ValueError("A maximum of 10 fish can be added to the aquarium")
    return {"tank_a": fish_list}


class TestAddFishToAquarium(unittest.TestCase):
    def test_add_fish_to_aquarium_success(self):
        actual = add_fish_to_aquarium(fish_list=["shark", "tuna"])
        expected = {"tank_a": ["rabbit"]}  # 错误的预期值
        self.assertEqual(actual, expected)

这个修改后的测试会失败，因为 add_fish_to_aquarium 不可能返回包含 "rabbit" 的鱼列表。

再次在同一目录下运行：

python -m unittest test_add_fish_to_aquarium.py

你会看到如下输出：

F
======================================================================
FAIL: test_add_fish_to_aquarium_success (test_add_fish_to_aquarium.TestAddFishToAquarium)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "test_add_fish_to_aquarium.py", line 13, in test_add_fish_to_aquarium_success
    self.assertEqual(actual, expected)
AssertionError: {'tank_a': ['shark', 'tuna']} != {'tank_a': ['rabbit']}
- {'tank_a': ['shark', 'tuna']}
+ {'tank_a': ['rabbit']}

----------------------------------------------------------------------
Ran 1 test in 0.001s

FAILED (failures=1)

失败输出表明测试未通过。实际输出 {'tank_a': ['shark', 'tuna']} 与我们在测试中设置的错误预期 {'tank_a': ['rabbit']} 不匹配。注意，输出第一行现在是 F 而不是 . —— . 表示通过，F 表示失败。

现在我们已经编写并运行了一个测试，接下来尝试为 add_fish_to_aquarium 函数的另一种行为编写另一个测试。

测试会抛出异常的函数

unittest 还可以帮助我们验证：当输入过多鱼类时，add_fish_to_aquarium 函数是否会正确抛出 ValueError 异常。我们在之前的例子基础上添加一个新的测试方法 test_add_fish_to_aquarium_exception：

test_add_fish_to_aquarium.py

import unittest

def add_fish_to_aquarium(fish_list):
    if len(fish_list) > 10:
        raise ValueError("A maximum of 10 fish can be added to the aquarium")
    return {"tank_a": fish_list}


class TestAddFishToAquarium(unittest.TestCase):
    def test_add_fish_to_aquarium_success(self):
        actual = add_fish_to_aquarium(fish_list=["shark", "tuna"])
        expected = {"tank_a": ["shark", "tuna"]}
        self.assertEqual(actual, expected)

    def test_add_fish_to_aquarium_exception(self):
        too_many_fish = ["shark"] * 25
        with self.assertRaises(ValueError) as exception_context:
            add_fish_to_aquarium(fish_list=too_many_fish)
        self.assertEqual(
            str(exception_context.exception),
            "A maximum of 10 fish can be added to the aquarium"
        )

新的测试方法 test_add_fish_to_aquarium_exception 向 add_fish_to_aquarium 传入一个包含 25 个 "shark" 的列表。

该测试使用 TestCase 提供的 with self.assertRaises(...) 上下文管理器来检查函数是否因输入过长而拒绝处理。self.assertRaises 的第一个参数是我们期望抛出的异常类（这里是 ValueError）。上下文管理器绑定到变量 exception_context，其 exception 属性包含函数实际抛出的 ValueError 对象。调用 str() 获取异常消息，验证其是否与预期一致。

在同一目录下运行：

python -m unittest test_add_fish_to_aquarium.py

输出如下：

..
----------------------------------------------------------------------
Ran 2 tests in 0.000s

OK

值得注意的是，如果 add_fish_to_aquarium 没有抛出异常，或者抛出了其他类型的异常（例如 TypeError 而非 ValueError），该测试就会失败。

提示：unittest.TestCase 除了 assertEqual 和 assertRaises 外，还提供了许多其他断言方法。完整列表请参考官方文档，以下是部分常用方法：

方法 断言条件
assertEqual(a, b) a == b
assertNotEqual(a, b) a != b
assertTrue(a) bool(a) is True
assertFalse(a) bool(a) is False
assertIsNone(a) a is None
assertIsNotNone(a) a is not None
assertIn(a, b) a in b
assertNotIn(a, b) a not in b

使用 setUp 方法创建资源

TestCase 还支持 setUp 方法，用于在每个测试前创建所需资源。当你有一段通用的准备代码需要在每个测试前运行时，setUp 可以避免重复。

看一个例子：

test_fish_tank.py

import unittest

class FishTank:
    def __init__(self):
        self.has_water = False

    def fill_with_water(self):
        self.has_water = True

class TestFishTank(unittest.TestCase):
    def setUp(self):
        self.fish_tank = FishTank()

    def test_fish_tank_empty_by_default(self):
        self.assertFalse(self.fish_tank.has_water)

    def test_fish_tank_can_be_filled(self):
        self.fish_tank.fill_with_water()
        self.assertTrue(self.fish_tank.has_water)

test_fish_tank.py 定义了一个 FishTank 类，其 has_water 初始为 False，调用 fill_with_water() 后变为 True。

TestFishTank 测试类中的 setUp 方法会在每个测试方法执行前创建一个新的 FishTank 实例，并赋值给 self.fish_tank。

因此，test_fish_tank_empty_by_default 和 test_fish_tank_can_be_filled 都会获得独立的 FishTank 实例。前者验证初始状态为无水，后者验证加水后状态变为有水。

在同一目录下运行：

python -m unittest test_fish_tank.py

输出：

..
----------------------------------------------------------------------
Ran 2 tests in 0.000s

OK

两个测试均通过。

提示：如果你有多个包含 TestCase 子类的测试文件，可以使用 python -m unittest discover 一次性运行所有测试。运行 python -m unittest discover --help 查看更多选项。

使用 tearDown 方法清理资源

TestCase 还提供了与 setUp 对应的 tearDown 方法，用于在每个测试结束后清理资源，例如关闭数据库连接或删除临时文件。

下面是一个使用 tearDown 清理文件系统的例子：

test_advanced_fish_tank.py

import os
import unittest

class AdvancedFishTank:
    def __init__(self):
        self.fish_tank_file_name = "fish_tank.txt"
        default_contents = "shark, tuna"
        with open(self.fish_tank_file_name, "w") as f:
            f.write(default_contents)

    def empty_tank(self):
        os.remove(self.fish_tank_file_name)


class TestAdvancedFishTank(unittest.TestCase):
    def setUp(self):
        self.fish_tank = AdvancedFishTank()

    def tearDown(self):
        self.fish_tank.empty_tank()

    def test_fish_tank_writes_file(self):
        with open(self.fish_tank.fish_tank_file_name) as f:
            contents = f.read()
        self.assertEqual(contents, "shark, tuna")

AdvancedFishTank 在初始化时创建 fish_tank.txt 文件并写入 "shark, tuna"，同时提供 empty_tank 方法删除该文件。

TestAdvancedFishTank 同时定义了 setUp 和 tearDown：

  • setUp 创建 AdvancedFishTank 实例；
  • tearDown 调用 empty_tank()，确保每次测试后删除文件，使每个测试都从干净状态开始。

test_fish_tank_writes_file 验证文件内容是否正确。

在同一目录下运行：

python -m unittest test_advanced_fish_tank.py

输出：

.
----------------------------------------------------------------------
Ran 1 test in 0.000s

OK

tearDown 允许你为 TestCase 子类中的所有测试编写统一的清理代码。

结论

在本教程中，你学习了：

  • 如何编写包含不同断言的 TestCase 类；
  • 如何使用 setUp 和 tearDown 方法；
  • 如何从命令行运行测试。

unittest 模块还提供了更多未在本教程中涵盖的类和工具。现在你已掌握基础知识，可以查阅 unittest 官方文档 了解更多信息。你可能还会对《如何为 Django 项目添加单元测试》感兴趣。