[FIXED] Testing cursor.execute with sql script called using pytest

Issue

Function to test

def get_adgroups_not_taked_share(
    campaign_ids: List[str], src_table: str, spend_src_table: str
) -> List[Tuple[str, str]]:


    loses_adgroups: List[Tuple[str, str]] = []

    with RedshiftCursor() as cursor:
        cursor.execute(
            """
            SELET some_data from some_table WHERE some_condition
            """
        )
        for row in cursor.fetchall():
            loses_adgroups.append((row[0], str(row[1])))

    return loses_adgroups

There is a test for this function

import pytest

from my_ap import get_adgroups_not_taked_share


@pytest.fixture
def campaigns_redshift_cursor_mock(mocker):
    cursor_mock = mocker.MagicMock()
    cursor_mock.fetchall.return_value = [
        ('hs_video544', '123123123', 100),
        ('hs_video547', '123123123', 50),
    ]

    rs_cursor_creator = mocker.patch('google_panel.logic.clean_creative.RedshiftCursor')
    rs_cursor_creator.return_value.__enter__.return_value = cursor_mock
    return rs_cursor_creator


@pytest.mark.django_db
def test_get_adgroups_not_taked_share(
        campaigns_redshift_cursor_mock,
        ):
    campaign_ids = ['1111', '2222', '3333']
    result = get_adgroups_not_taked_share(campaign_ids, 'test_table', 'spend_src_table')
    assert result == [('hs_video544', '123123123'), ('hs_video547', '123123123')]

Now I want to add a new feature to test the sql script. Checking that the correct sql query is being called. something like

def test_get_adgroups_not_taked_share(
            campaigns_redshift_cursor_mock,
            ):
    ......
    query = """SELET some_data from some_table WHERE some_condition"""
    campaigns_redshift_cursor_mock.execute.assert_called_with(query)

But got

E       AssertionError: Expected call: execute('query')
E       Not called

Solution

The short answer is that you need to assert this: campaigns_redshift_cursor_mock.return_value.__enter__.return_value.execute.assert_called_once_with(query). The reason is that you're using RedshiftCursor as a context manager (hence the return_value.__enter__.return_value part) and only then calling the execute method.

A slightly longer answer is how I got this assert.

I wrote this library which adds the mg pytest fixture. To use it, pip install it and then add the mg fixture to your test and execute its generate_asserts method like so:

@pytest.mark.django_db
def test_get_adgroups_not_taked_share(
        campaigns_redshift_cursor_mock,
        mg
        ):
    campaign_ids = ['1111', '2222', '3333']
    result = get_adgroups_not_taked_share(campaign_ids, 'test_table', 'spend_src_table')
    mg.generate_asserts(campaigns_redshift_cursor_mock)
    assert result == [('hs_video544', '123123123'), ('hs_video547', '123123123')]

Then run the test as usual and you would get this output to the console:

assert 1 == campaigns_redshift_cursor_mock.call_count
campaigns_redshift_cursor_mock.assert_called_once_with()
campaigns_redshift_cursor_mock.return_value.__enter__.assert_called_once_with()
campaigns_redshift_cursor_mock.return_value.__enter__.return_value.execute.assert_called_once_with('\n            SELET some_data from some_table WHERE some_condition\n            ')
campaigns_redshift_cursor_mock.return_value.__enter__.return_value.fetchall.assert_called_once_with()
campaigns_redshift_cursor_mock.return_value.__exit__.assert_called_once_with(None, None, None)

These are all the calls made to the mock, you can filter the ones which are relevant for you.

Answered By - Peter K