时间戳

首页 > 解决方案 > 如何在 Python 3 单元测试中模拟“with connect”SQL 查询？

python - 如何在 Python 3 单元测试中模拟“with connect”SQL 查询？

问题描述

对于像下面这样的函数，其中引擎是 SQLAlchemy 引擎，例如，

from sqlalchemy import create_engine
engine = create_engine(f'mysql+pymysql://{db_username}:{db_password}@{db_host}:{db_port}/{db_name}', pool_recycle=3600)

def pull(engine, afield):
    query = f"SELECT col1, col2 FROM mydatabase WHERE field='{afield}'"
    with engine.connect() as conn:
        result = conn.execute(query)
    return result

如何模拟单元测试中“with”语句的查询结果？

标签： pythonunit-testingsqlalchemymocking

解决方案

假设上述函数datapull.py位于与本单元测试代码命名并位于同一目录的文件中，那么下面是一个模拟 SQL 结果的示例。

模拟引擎和连接相当简单，但模拟引擎返回连接则比较棘手。

with 语句（此处__enter__针对 Python 3 进行了描述）在引擎上调用后返回连接。

因此mock_engine.connect.return_value.__enter__.return_value = mock_conn将在上述代码的 with 语句中正确返回模拟连接。

import mock
import unittest
import datapull


class TestQueryRDS(unittest.TestCase):
    def test_happy_day(self):
        """Test happy day query"""
        mock_conn = mock.MagicMock()
        mock_conn.execute.return_value = [['c1', 'c2'], ['a1', 'a2']]
        mock_engine = mock.MagicMock()
        mock_engine.connect.return_value.__enter__.return_value = mock_conn
        actual = datapull.pull(mock_engine, '1234abcd')
        assert actual == [['c1', 'c2'], ['a1', 'a2']]
        exp_sql = "SELECT col1, col2 FROM mydatabase WHERE field='1234abcd'"
        mock_conn.execute.assert_called_once_with(exp_sql)

推荐阅读