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

【问题标题】：How do I mock a "with connect" SQL query in a Python 3 unit testing?如何在 Python 3 单元测试中模拟“with connect”SQL 查询？
【发布时间】：2020-05-04 23:46:46
【问题描述】：

对于像下面这样的函数，其中引擎是 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”语句的查询结果？

【问题讨论】：

标签： python unit-testing sqlalchemy mocking

【解决方案1】：

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

模拟引擎和连接相当简单，但模拟引擎以返回连接更棘手。

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

因此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)

【讨论】：