如何在 Python 3 單元測試中模擬“with connect”SQL 查詢？

Question

對於如下所示的 function，其中引擎是 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”語句的查詢結果？

Answer 1

假設上面的 function 在一個名為datapull.py的文件中並且與這個單元測試代碼位於同一目錄中，那么下面是 mocking SQL 結果的示例。

Mocking 引擎和連接相當簡單，但是 mocking返回連接的引擎比較棘手。

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)