3、单元测试基础：Python unittest框架、pytest入门、Mock对象与桩模块

单元测试，说白了就是给代码做「体检」。你写了一个函数，它能不能跑？边界情况处理了没？异常抛得对不对？这些都得靠单元测试来兜底。我在做量化交易系统的时候，最怕的就是策略逻辑里藏了个隐蔽的 bug，回测跑得漂漂亮亮，实盘一上线就亏钱。嗯，从那以后，我养成了一个习惯：先写测试，再写代码。

这一章，咱们就聊聊 Python 里最常用的两个测试框架——unittest 和 pytest，以及怎么用 Mock 对象和桩模块来隔离外部依赖。你想想看，如果你的策略代码里调用了行情接口，总不能每次跑测试都真的去连交易所吧？

3.1 unittest 框架：Python 自带的「老大哥」

unittest 是 Python 标准库的一部分，开箱即用。它的设计思路借鉴了 Java 的 JUnit，所以如果你写过 Java 测试，上手会非常快。

我个人习惯用 unittest 来写一些比较「正统」的测试用例，尤其是团队里有人对第三方库有顾虑的时候。毕竟它是内置的，零依赖。

3.1.1 基本用法

一个最简单的测试用例，需要继承 unittest.TestCase，然后定义以 test_ 开头的方法。每个方法就是一个独立的测试。

import unittest

class TestOrderManager(unittest.TestCase):

    def test_create_limit_order(self):
        # 模拟一个限价单创建逻辑
        order = {'symbol': 'BTCUSDT', 'side': 'buy', 'price': 50000, 'quantity': 0.1}
        self.assertEqual(order['price'], 50000)
        self.assertIn('symbol', order)

    def test_invalid_quantity(self):
        # 测试数量为负数时是否抛出异常
        with self.assertRaises(ValueError):
            create_order(quantity=-1)  # 假设这个函数会校验

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

这里有几个常用的断言方法，我列个表给你参考：

断言方法 用途
assertEqual(a, b) 检查 a == b
assertTrue(x) 检查 x 为 True
assertIn(a, b) 检查 a 在 b 中
assertRaises(Exc, fun) 检查是否抛出指定异常
assertAlmostEqual(a, b) 浮点数近似相等（精度问题）
小提示： 浮点数比较是个坑。我在做回测引擎时，遇到过因为浮点精度导致测试失败的情况。后来统一用 assertAlmostEqual 或者 round() 处理，才消停。

3.1.2 setUp 与 tearDown

每个测试方法执行前后，都可以通过 setUp 和 tearDown 来做一些准备工作或清理工作。比如初始化一个数据库连接，或者创建临时文件。

class TestDatabase(unittest.TestCase):

    def setUp(self):
        # 每个测试前执行
        self.db = create_test_database()

    def tearDown(self):
        # 每个测试后执行
        self.db.close()

    def test_insert(self):
        self.db.insert('order', {'id': 1})
        result = self.db.query('SELECT * FROM orders')
        self.assertEqual(len(result), 1)

如果你希望所有测试只执行一次 setUp，可以用 @classmethod 装饰的 setUpClass。这个在测试耗时操作时特别有用。

3.2 pytest：更简洁、更强大

说实话，我后来慢慢转向了 pytest。为什么？因为它写起来更「Pythonic」。你不需要继承任何类，不需要写 self.assertXxx，直接用 assert 就行。

3.2.1 安装与入门

安装很简单：

pip install pytest

然后写一个测试文件，比如 test_strategy.py：

def test_calculate_sma():
    prices = [100, 102, 101, 103, 105]
    sma = sum(prices[-3:]) / 3
    assert sma == 103.0  # 直接用 assert

在终端运行 pytest，它会自动发现所有以 test_ 开头的函数或文件。输出结果也很漂亮，绿色通过，红色失败。

3.2.2 Fixture：依赖注入的利器

Fixture 是 pytest 最让我喜欢的功能之一。它相当于 unittest 里的 setUp，但更灵活。你可以把 fixture 当作一个「可复用的测试资源」。

import pytest

@pytest.fixture
def order_data():
    # 返回一个测试用的订单数据
    return {'symbol': 'ETHUSDT', 'price': 3000, 'quantity': 0.5}

def test_order_price(order_data):
    assert order_data['price'] == 3000

def test_order_quantity(order_data):
    assert order_data['quantity'] == 0.5

你看，order_data 这个 fixture 被自动注入到了测试函数里。不用手动创建，不用清理，pytest 会帮你管理生命周期。

核心要点： Fixture 的作用域可以设置为 function、class、module 或 session。如果你有一个耗时的初始化操作（比如加载模型），建议用 scope="session"，这样整个测试过程只执行一次。

3.2.3 参数化测试

有时候你需要用多组数据测试同一个逻辑。用 unittest 你得写多个方法，或者用循环。pytest 提供了 @pytest.mark.parametrize，一行搞定。

@pytest.mark.parametrize("input,expected", [
    (100, 100.0),
    (99.99, 99.99),
    (0, 0.0),
    (-1, ValueError),  # 期望抛出异常
])
def test_validate_price(input, expected):
    if expected == ValueError:
        with pytest.raises(ValueError):
            validate_price(input)
    else:
        assert validate_price(input) == expected

这样做的好处是，如果某组数据失败了，pytest 会明确告诉你「第3组数据失败了」，而不是笼统地说「测试失败」。

3.3 Mock对象与桩模块：隔离外部依赖

做量化测试，最头疼的就是外部依赖。行情数据、交易接口、数据库……这些在测试环境里往往不可用，或者不稳定。这时候就需要 Mock 和 Stub 出场了。

简单区分一下：

  • Stub（桩）：提供一个固定的返回值，比如模拟行情接口返回一个固定的价格。
  • Mock（模拟对象）：不仅能返回固定值，还能记录调用次数、参数，甚至验证行为。

Python 的 unittest.mock 模块（Python 3.3+ 内置）提供了强大的 Mock 功能。pytest 也可以配合 pytest-mock 插件使用。

3.3.1 使用 Mock 替换外部调用

假设你的策略里调用了行情接口：

# strategy.py
import market_data

def should_buy(symbol):
    price = market_data.get_price(symbol)
    return price < 50000

测试时，你不能真的去连交易所。用 Mock 把它替换掉：

from unittest.mock import patch
import strategy

@patch('strategy.market_data.get_price')
def test_should_buy_when_price_low(mock_get_price):
    mock_get_price.return_value = 48000
    result = strategy.should_buy('BTCUSDT')
    assert result is True
    mock_get_price.assert_called_once_with('BTCUSDT')

这里 @patch 装饰器把 market_data.get_price 替换成了一个 Mock 对象。我们指定它返回 48000，然后验证策略逻辑是否正确。最后还检查了 get_price 是否被正确调用了一次。

注意： @patch 的路径要写对。我曾经因为路径写错，Mock 没生效，测试跑了一晚上才发现。建议用 @patch('模块名.函数名') 这种完整路径。

3.3.2 用 Stub 模拟复杂对象

有时候你需要模拟一个完整的对象，比如一个数据库连接或者一个 WebSocket 客户端。这时候可以自己写一个简单的 Stub 类。

class StubDatabase:
    def __init__(self):
        self.data = {}

    def insert(self, table, record):
        if table not in self.data:
            self.data[table] = []
        self.data[table].append(record)

    def query(self, sql):
        # 简单模拟，实际项目里可以更复杂
        return self.data.get('orders', [])

def test_save_order():
    db = StubDatabase()
    order = {'id': 1, 'symbol': 'BTC'}
    save_order_to_db(db, order)  # 假设这个函数会调用 db.insert
    assert len(db.query('SELECT * FROM orders')) == 1

Stub 的好处是轻量、可控。你不需要引入任何第三方库，自己写一个简单的类就行。我在做回测系统时，经常用 Stub 来模拟交易所的订单簿和账户余额。

3.3.3 Mock 的高级用法：验证行为

Mock 不仅能返回值，还能记录调用历史。这在测试「某个函数是否被正确调用」时特别有用。

from unittest.mock import MagicMock

def test_order_execution():
    exchange = MagicMock()
    exchange.place_order.return_value = {'order_id': '12345'}

    result = execute_trade(exchange, 'buy', 100)
    assert result['order_id'] == '12345'
    exchange.place_order.assert_called_once_with(
        side='buy', quantity=100
    )

你看，MagicMock 会自动创建所有属性和方法。你只需要设置 return_value，然后调用 assert_called_once_with 来验证参数。如果调用次数不对，或者参数不匹配，测试会直接失败。

避坑指南： 我曾经在测试一个异步回调函数时，忘了 Mock 的 side_effect 参数。它可以让 Mock 在每次调用时返回不同的值，或者抛出异常。比如 mock.get_price.side_effect = [100, 200, ConnectionError]，这样第一次返回100，第二次返回200，第三次抛出连接异常。非常适合测试重试逻辑。

3.4 本章小结

单元测试不是「写完了再补」的事，而是开发过程中的一部分。我个人习惯在写每个函数之前，先想清楚它的输入输出和边界条件，然后写一个测试用例。这样代码写起来反而更快，因为你知道它「应该」怎么工作。

这一章我们聊了：

  • unittest：Python 内置，适合传统项目，结构清晰但略显啰嗦。
  • pytest：更简洁，fixture 和参数化测试是亮点，推荐新项目使用。
  • Mock 与 Stub：隔离外部依赖，让测试变得可控、可重复。

记住一句话：没有单元测试的代码，就像没有安全带的赛车。 跑得快，但翻车的时候也最惨。下一章我们会深入测试覆盖率，看看你的测试到底「测够了没」。

专注资料整理
← 上一章 📖 返回目录 下一章 →