如何在 `with` 上下文管理器中模拟 sqlalchemy 执行输出

Question

我正在尝试为我的 python 脚本编写单元测试，该脚本使用 sqlalchemy 连接到 MySQL。

我的函数如下所示：

def check_if_table_exists(db):
    with db.connect() as cursor:
        table_exists = cursor.execute(f"SHOW TABLES LIKE '{PRIMARY_TABLE_NAME}';")
        if not table_exists.rowcount:
            cursor.execute(f"CREATE TABLE...

我找不到任何关于如何首先模拟 db.connect() 的资源，这反过来也需要对其 execute 进行模拟，以便我可以测试不同的 table_exists 场景.也有可能我的代码根本无法与适当的单元测试相结合，我需要使用游标对象调用该函数。

作为参考，db 是 sqlalchemy.create_engine 的输出。

TLDR 对于我为 SHOW 语句取回行和没有取回行的情况，我需要帮助开始进行单元测试。

Answer 1

要修补上下文管理器，您必须修补 __enter__ 的 return 值，它会在进入上下文管理器时调用。这是您的代码示例：

from unittest import mock
from sqlalchemy import create_engine
from my_project.db_connect import check_if_table_exists


@mock.patch('sqlalchemy.engine.Engine.connect')
def test_dbconnect(engine_mock):
    db = create_engine('sqlite:///:memory:')
    cursor_mock = engine_mock.return_value.__enter__.return_value
    cursor_mock.execute.return_value.rowcount = 0
    check_if_table_exists(db)
    cursor_mock.execute.assert_called_with("CREATE TABLE")

在此代码中，engine_mock.return_value 是 Engine 的模拟实例，要获取 cursor 的模拟实例，您需要按照说明添加 __enter__.return_value。

有了这个，您现在可以模拟 execute 的 return 值 - 在这种情况下，您只对代码中选中的 rowcount 属性感兴趣。请注意，这将更改 execute 的所有调用的 return 值 - 如果后续调用需要不同的值，则可以改用 side_effect。

如何在 `with` 上下文管理器中模拟 sqlalchemy 执行输出

How to mock sqlalchemy execution output inside `with` context manager

python

testing

unit-testing

sqlalchemy

pymysql