CODEWITHSUNDEEP
pythonunit-testingsqlalchemymocking
pettinato
pettinato

Reputation: 1542

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

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?

Upvotes: 1

Views: 1856

Answers (1)

pettinato
pettinato

Reputation: 1542

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)

Upvotes: 1

Related Questions