如何对包含Beanie查询的Python代码进行单元测试与Mock

嘿，我完全理解你现在的困扰——要测试用了Beanie（MongoDB的异步ODM）的代码，又不想真的连接数据库对吧？其实核心就是把Record.find_one这个异步查询给Mock掉，分两种场景验证逻辑就好，我给你一步步拆解：

首先先理清楚你的代码结构：

业务代码（core_code.py） #

from mongo_models import Record

async def do_something(record_id: str):
    record = await Record.find_one(Record.id == record_id)
    if record is None:
        record = Record(id='randomly_generated_string', content='')
    # 其他操作忽略
    return record

模型定义（mongo_models.py） #

from beanie import Document

class Record(Document):
    id: str
    content: str

你之前尝试用AsyncMock模拟Record实例，但方向有点偏——我们需要Mock的是Record.find_one这个异步类方法，而不是单个实例。接下来给你完整的测试方案：

第一步：准备测试依赖 #

确保你已经安装了必要的测试工具：

pip install pytest pytest-asyncio pytest-mock

第二步：编写分场景的测试用例 #

我们需要覆盖两种核心场景：记录存在时的逻辑，和记录不存在时创建新记录的逻辑。

import pytest
from unittest.mock import AsyncMock
from core_code import do_something
from mongo_models import Record

@pytest.mark.asyncio
async def test_do_something_when_record_exists(mocker):
    """测试：当数据库中存在对应记录时，函数返回该记录"""
    # 1. 模拟一个存在的Record实例
    test_record = Record(id="test_123", content="hello world")
    
    # 2. Mock业务代码中用到的Record.find_one异步方法
    # 注意路径是core_code.Record.find_one，要Mock业务代码实际使用的类方法
    mock_find_one = mocker.patch("core_code.Record.find_one", new_callable=AsyncMock)
    mock_find_one.return_value = test_record

    # 3. 调用测试函数
    result = await do_something("test_123")

    # 4. 验证逻辑正确性
    mock_find_one.assert_awaited_once_with(Record.id == "test_123")
    assert result == test_record

@pytest.mark.asyncio
async def test_do_something_when_record_not_exists(mocker):
    """测试：当数据库中没有对应记录时，函数创建新记录并返回"""
    # 1. Mock Record.find_one返回None，表示记录不存在
    mock_find_one = mocker.patch("core_code.Record.find_one", new_callable=AsyncMock)
    mock_find_one.return_value = None

    # 2. 调用测试函数
    result = await do_something("non_existent_id")

    # 3. 验证逻辑正确性
    mock_find_one.assert_awaited_once_with(Record.id == "non_existent_id")
    assert isinstance(result, Record)
    assert result.content == ""
    assert result.id != "non_existent_id"  # 确认是新生成的随机ID

关键注意点 #

• Mock路径要准确：必须Mock业务代码（core_code）里导入的Record的find_one方法，而不是原始的mongo_models.Record.find_one——Python的导入是按模块引用的，我们要替换的是业务代码实际使用的类方法。
• 用AsyncMock处理异步方法：因为find_one是异步函数，必须用AsyncMock模拟，这样才能用assert_awaited_once_with这类异步断言验证调用情况。
• 覆盖所有分支：两种场景都要测试，确保函数逻辑没有遗漏。

这样跑测试就完全不需要连接真实的MongoDB，纯单元测试就能验证你的业务逻辑啦！

备注：内容来源于stack exchange，提问作者Nordico