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

Question

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

我的函数如下所示：

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__，该值在进入上下文管理器时被调用。这是您的代码的示例：

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感兴趣。rowcount请注意，这将更改所有调用的返回值execute- 如果后续调用需要不同的值，则可以改用side_effect。