How do I mock a "with connect" SQL query in a Python 3 unit testing?

Question

For a function like the below, where engine is a SQLAlchemy engine like,

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

How do I mock the query result from the "with" statement in a unit test?

Answer 1

Assuming the above function is in a file named datapull.py and located in the same directory as this unit test code, then the below is an example of mocking the SQL result.

Mocking the engine and connection are fairly simple, but mocking the engine to return the connection is trickier.

The with statement (described for Python 3 here ) returns the connection after __enter__ is called on the engine.

Thus mock_engine.connect.return_value.__enter__.return_value = mock_conn will properly return the mock connection in the above code's with statement.

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)