在Python的单元测试场景中,我们常常需要测试某个函数或类的逻辑,但这些逻辑可能依赖外部API、数据库查询、文件读写等难以在测试环境中直接调用的组件。这时候使用mock来模拟这些依赖对象,就可以让测试脱离外部环境的限制,更精准地验证核心逻辑的正确性。

Python mock基础概念
mock是Python标准库unittest的一部分,位于unittest.mock模块中,主要作用是创建一个模拟对象,替代测试中的真实依赖。模拟对象可以自定义返回值、抛出异常,还可以记录自身的调用情况,方便我们验证代码是否按预期调用了依赖。
最常用的mock类是Mock和MagicMock,后者是前者的子类,支持更多魔术方法的模拟,日常使用中如果没有特殊需求,优先选择MagicMock即可。
基本使用:模拟函数返回值
假设我们有一个函数get_user_info,它依赖外部接口获取用户数据,我们想要测试这个函数的处理逻辑,不需要真的调用外部接口,只需要模拟接口的返回值即可。
首先看被测试的代码:
# 被测试的代码模块 user_service.py
import requests
def get_user_info(user_id):
# 调用外部接口获取用户信息
resp = requests.get(f"http://ipipp.com/api/user/{user_id}")
if resp.status_code == 200:
return resp.json()
return None
def process_user_data(user_id):
# 处理用户数据的业务逻辑
user_info = get_user_info(user_id)
if user_info is None:
return "用户不存在"
if user_info.get("age") > 18:
return "成年用户"
return "未成年用户"
接下来编写测试代码,模拟requests.get的返回值:
# 测试代码 test_user_service.py
import unittest
from unittest.mock import patch, MagicMock
from user_service import process_user_data
class TestProcessUserData(unittest.TestCase):
# 使用patch装饰器模拟requests.get方法
@patch("user_service.requests.get")
def test_adult_user(self, mock_get):
# 创建模拟的响应对象
mock_resp = MagicMock()
mock_resp.status_code = 200
mock_resp.json.return_value = {"name": "张三", "age": 20}
# 设置mock_get的返回值为模拟响应
mock_get.return_value = mock_resp
result = process_user_data(1)
# 验证结果
self.assertEqual(result, "成年用户")
# 验证requests.get是否被正确调用
mock_get.assert_called_once_with("http://ipipp.com/api/user/1")
@patch("user_service.requests.get")
def test_user_not_exist(self, mock_get):
mock_resp = MagicMock()
mock_resp.status_code = 404
mock_get.return_value = mock_resp
result = process_user_data(2)
self.assertEqual(result, "用户不存在")
if __name__ == "__main__":
unittest.main()
模拟类实例和方法
除了模拟函数,我们还经常需要模拟类的实例。比如有一个数据库操作类,我们想要模拟它的查询方法,避免真的连接数据库。
被测试的代码:
# db_service.py
class UserDB:
def __init__(self, db_url):
self.db_url = db_url
def query_user(self, user_id):
# 实际会连接数据库查询,这里简化为返回固定值
return {"id": user_id, "name": "默认用户"}
def get_user_name(db, user_id):
user = db.query_user(user_id)
return user.get("name")
测试代码模拟UserDB的实例:
# test_db_service.py
import unittest
from unittest.mock import MagicMock
from db_service import get_user_name, UserDB
class TestGetUserName(unittest.TestCase):
def test_get_user_name(self):
# 创建模拟的UserDB实例
mock_db = MagicMock(spec=UserDB)
# 设置query_user方法的返回值
mock_db.query_user.return_value = {"id": 1, "name": "李四"}
result = get_user_name(mock_db, 1)
self.assertEqual(result, "李四")
# 验证query_user方法被调用
mock_db.query_user.assert_called_once_with(1)
if __name__ == "__main__":
unittest.main()
实用测试技巧
1. 使用spec参数限制模拟对象的行为
创建Mock或MagicMock时,传入spec参数可以指定模拟对象需要符合的真实对象的接口,避免模拟对象被调用不存在的方法时不会报错,提升测试的准确性。
from unittest.mock import MagicMock from db_service import UserDB # 模拟UserDB实例,限制只能调用真实存在的方法 mock_db = MagicMock(spec=UserDB) # 调用真实存在的方法不会报错 mock_db.query_user(1) # 调用不存在的方法会直接抛出AttributeError # mock_db.undefined_method() # 这行会报错
2. 模拟抛出异常
如果需要测试代码在依赖抛出异常时的处理逻辑,可以通过side_effect参数设置模拟对象抛出异常。
import unittest
from unittest.mock import patch, MagicMock
from user_service import process_user_data
class TestProcessUserDataException(unittest.TestCase):
@patch("user_service.requests.get")
def test_request_exception(self, mock_get):
# 设置调用requests.get时抛出连接异常
mock_get.side_effect = ConnectionError("连接失败")
result = process_user_data(1)
# 这里可以根据实际代码的异常处理逻辑验证结果
self.assertEqual(result, "用户不存在")
if __name__ == "__main__":
unittest.main()
3. 验证调用次数和参数
mock对象提供了多个断言方法,用来验证依赖的调用情况:
assert_called_once():验证只被调用一次assert_called_once_with(*args, **kwargs):验证只被调用一次且参数匹配assert_called_with(*args, **kwargs):验证最后一次调用的参数匹配assert_has_calls(calls_list):验证调用序列匹配
from unittest.mock import MagicMock mock_func = MagicMock() mock_func(1, 2, name="test") mock_func(3) # 验证调用次数 self.assertTrue(mock_func.call_count == 2) # 验证最后一次调用参数 mock_func.assert_called_with(3) # 验证第一次调用参数 mock_func.assert_any_call(1, 2, name="test")
4. 使用patch.object模拟对象的方法
如果只需要模拟某个对象的单个方法,不需要模拟整个对象,可以使用patch.object。
import unittest
from unittest.mock import patch, MagicMock
from db_service import UserDB, get_user_name
class TestPatchObject(unittest.TestCase):
def test_patch_object(self):
db = UserDB("http://ipipp.com/db")
# 只模拟db实例的query_user方法
with patch.object(db, "query_user") as mock_query:
mock_query.return_value = {"id": 1, "name": "王五"}
result = get_user_name(db, 1)
self.assertEqual(result, "王五")
if __name__ == "__main__":
unittest.main()
注意事项
使用mock时需要注意模拟的路径,patch装饰器的第一个参数是要模拟的对象所在的命名空间路径,而不是定义该对象的路径。比如要模拟user_service模块中使用的requests.get,路径应该是user_service.requests.get,而不是requests.get,否则可能无法生效。
另外,不要过度使用mock,只有当依赖确实难以在测试环境中调用,或者调用成本很高时才使用mock。如果依赖本身很容易构造,直接使用真实对象测试会更可靠。