How to mock sqlalchemy execution output inside `with` context manager

Question

I'm trying to write unit tests for my python script which uses sqlalchemy to connect to MySQL.

My function looks like this:

def check_if_table_exists(db):
    with db.connect() as cursor:
        table_exists = cursor.execute(f"SHOW TABLES LIKE '{PRIMARY_TABLE_NAME}';")
        if not table_exists.rowcount:
            cursor.execute(f"CREATE TABLE...

I can't find any resources on how to first mock out db.connect(), which in turn also needs to have its execute mocked so that I can test different table_exists scenarios. It's also possible that my code is simply not cohesive to proper unit testing and I need to call the function with a cursor object to begin with.

For reference, db is the output of sqlalchemy.create_engine.

TLDR I need help getting started on unit testing for cases where I get rows back for the SHOW statement and when I don't.

For mocking a context manager check out [this answer](https://stackoverflow.com/a/28852060/12480730). — MrBean Bremen, Jun 02 '20 at 09:13

score 6 · Accepted Answer · answered Jun 02 '20 at 18:26

To patch a context manager, you have to patch the return value of __enter__, which is called on entering a context manager. Here is an example for your code:

from unittest import mock
from sqlalchemy import create_engine
from my_project.db_connect import check_if_table_exists


@mock.patch('sqlalchemy.engine.Engine.connect')
def test_dbconnect(engine_mock):
    db = create_engine('sqlite:///:memory:')
    cursor_mock = engine_mock.return_value.__enter__.return_value
    cursor_mock.execute.return_value.rowcount = 0
    check_if_table_exists(db)
    cursor_mock.execute.assert_called_with("CREATE TABLE")

In this code engine_mock.return_value is a mocked instance of Engine, and to get the mock for cursor, you need to add __enter__.return_value as described.

Having this, you can now mock the return value of execute - in this case you are only interested in the rowcount attribute which is checked in the code. Note that this will change the return value for all calls of execute - if you need different values for subsequent calls, you can use side_effect instead.

i prefer the `with mock.patch("sqlalchemy.engine.Engine.connect") as engine_mock:` syntax but otherwise this worked like a charm, thanks! — notacorn, Jun 03 '20 at 02:09

How to mock sqlalchemy execution output inside `with` context manager

1 Answers1