Skip to content

Python 单元测试中的 `unittest.mock` 应用详解

Posted on:2024年11月23日 at 13:10

在软件开发过程中,单元测试是一项至关重要的实践。为了确保每个模块和函数按预期工作,开发者常常需要模拟(Mock)外部依赖,例如数据库连接、网络请求等。Python 的 unittest.mock 库为这一需求提供了强大的支持。本篇文章将深入探讨如何在 Python 单元测试中使用 unittest.mock,涵盖基本用法、进阶技巧以及实际应用案例。

基本概念

unittest.mock 提供了多个类和函数,最常用的包括:

创建 Mock 对象

创建一个基本的 Mock 对象非常简单:

from unittest.mock import Mock

mock = Mock()
mock.some_method.return_value = "mocked!"
result = mock.some_method()
print(result)  # 输出: mocked!

在上述示例中,some_method 是一个动态创建的方法,其返回值被设置为 "mocked!"

配置 Mock 的返回值和副作用

返回值

可以通过 return_value 属性设置 Mock 对象被调用时的返回值:

mock = Mock(return_value=10)
print(mock())  # 输出: 10

副作用

side_effect 属性允许你在 Mock 对象被调用时执行特定的行为,例如抛出异常或返回不同的值:

from unittest.mock import Mock

# 抛出异常
mock = Mock(side_effect=ValueError("An error occurred"))
try:
    mock()
except ValueError as e:
    print(e)  # 输出: An error occurred

# 返回不同的值
mock = Mock(side_effect=[1, 2, 3])
print(mock(), mock(), mock())  # 输出: 1 2 3

断言 Mock 的调用

unittest.mock 提供了多种断言方法,用于验证 Mock 对象是否按预期被调用。

断言至少被调用一次

mock = Mock()
mock()
mock.assert_called()

断言只被调用一次

mock = Mock()
mock()
mock.assert_called_once()

断言被调用时的参数

mock = Mock()
mock(1, 2, key='value')
mock.assert_called_with(1, 2, key='value')

断言任意一次被特定参数调用

mock = Mock()
mock(1)
mock(2, key='value')
mock.assert_any_call(2, key='value')

断言没有被调用

mock = Mock()
mock.assert_not_called()

使用 patch 装饰器和上下文管理器

patchunittest.mock 中最常用的功能之一,用于临时替换模块或类中的对象。

作为装饰器使用

from unittest.mock import patch

@patch('module.ClassName')
def test(mock_class):
    instance = mock_class.return_value
    instance.method.return_value = 'mocked!'
    assert module.ClassName().method() == 'mocked!'

在这个例子中,module.ClassName 被替换为一个 Mock 对象,mock_class 被传递给测试函数以供配置和断言。

作为上下文管理器使用

from unittest.mock import patch

with patch('module.ClassName') as mock_class:
    instance = mock_class.return_value
    instance.method.return_value = 'mocked!'
    assert module.ClassName().method() == 'mocked!'

使用 with 语句可以确保 Mock 对象只在特定代码块内生效,退出代码块后自动恢复原状。

自动规格化(Autospeccing)

自动规格化使得 Mock 对象模拟的对象具有真实对象的接口,这可以帮助捕捉代码中的 API 变更或错误使用。

使用 autospec=True

from unittest.mock import patch

@patch('module.ClassName', autospec=True)
def test(mock_class):
    instance = mock_class.return_value
    instance.method.return_value = 'mocked!'
    assert module.ClassName().method() == 'mocked!'

通过设置 autospec=True,Mock 对象会根据被替换的 ClassName 自动生成接口,任何不符合接口的调用都会抛出 AttributeError

使用 create_autospec

from unittest.mock import create_autospec

def some_function(a, b):
    pass

mock_function = create_autospec(some_function, return_value='mocked!')
print(mock_function(1, 2))  # 输出: mocked!

create_autospec 函数允许你基于任意对象创建规格化的 Mock 对象。

支持魔术方法

MagicMock 类预配置了大多数魔术方法(如 __str__, __iter__ 等),使其能够模拟实现 Python 协议的对象。

示例:模拟迭代器

from unittest.mock import MagicMock

mock = MagicMock()
mock.__iter__.return_value = iter([1, 2, 3])

for item in mock:
    print(item)  # 输出: 1 2 3

示例:模拟上下文管理器

from unittest.mock import MagicMock

mock = MagicMock()
mock.__enter__.return_value = 'resource'
mock.__exit__.return_value = False

with mock as resource:
    print(resource)  # 输出: resource

辅助工具

sentinel

sentinel 提供了一组唯一的对象,适用于需要在测试中比较对象身份的场景。

from unittest.mock import sentinel

mock = Mock()
mock.return_value = sentinel.result

result = mock()
assert result is sentinel.result

ANY

ANY 可以用于匹配任何值,用于在断言中忽略特定参数。

from unittest.mock import Mock, ANY

mock = Mock()
mock(1, key='value')

mock.assert_called_with(1, key=ANY)

实战案例

假设我们有一个模块 weather.py,其中有一个函数 get_weather 依赖于外部 API 获取天气信息。为了单元测试 get_weather,我们可以使用 unittest.mock 来模拟 API 的响应。

weather.py

import requests

def get_weather(city):
    response = requests.get(f'http://api.weather.com/{city}')
    if response.status_code == 200:
        return response.json()['weather']
    else:
        return 'Unknown'

test_weather.py

from unittest import TestCase
from unittest.mock import patch
import weather

class TestWeather(TestCase):
    
    @patch('weather.requests.get')
    def test_get_weather_success(self, mock_get):
        mock_response = mock_get.return_value
        mock_response.status_code = 200
        mock_response.json.return_value = {'weather': 'Sunny'}
        
        result = weather.get_weather('Beijing')
        self.assertEqual(result, 'Sunny')
        mock_get.assert_called_with('http://api.weather.com/Beijing')
    
    @patch('weather.requests.get')
    def test_get_weather_failure(self, mock_get):
        mock_response = mock_get.return_value
        mock_response.status_code = 404
        
        result = weather.get_weather('UnknownCity')
        self.assertEqual(result, 'Unknown')
        mock_get.assert_called_with('http://api.weather.com/UnknownCity')

在这个测试案例中,我们使用 patch 装饰器替换了 requests.get 方法,模拟了不同的返回状态码和响应内容,从而测试了 get_weather 函数在不同情况下的行为。

总结

unittest.mock 是 Python 单元测试中不可或缺的工具,通过模拟外部依赖,开发者可以编写高效、隔离的测试用例。本文介绍了 unittest.mock 的基本用法、进阶技巧以及实际应用,涵盖了 Mock 对象的创建、配置、断言以及更复杂的场景如自动规格化和魔术方法的支持。 掌握 unittest.mock 不仅能够提升测试的覆盖率和可靠性,还能帮助开发者更好地理解和设计代码的依赖关系。建议读者在实际项目中多加练习,熟练运用 unittest.mock 来构建强大的测试套件。