使用 pytest 調用的 sql 腳本測試 cursor.execute

Question

Function來測試

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

這個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')]

現在我想添加一個新功能來測試 sql 腳本。 檢查是否正在調用正確的 sql 查詢。 就像是

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)

但是得到了

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

Answer 1

簡短的回答是您需要斷言： campaigns_redshift_cursor_mock.return_value.__enter__.return_value.execute.assert_called_once_with(query) 。 原因是您使用RedshiftCursor作為上下文管理器（因此是return_value.__enter__.return_value部分），然后才調用execute方法。

一個稍長的答案是我如何得到這個斷言。

我編寫了這個庫，它添加了mg pytest 夾具。 要使用它，pip 安裝它，然后將mg夾具添加到您的測試並執行它的generate_asserts方法，如下所示：

@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')]

然后像往常一樣運行測試，你會得到這個 output 到控制台：

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)

這些是對模擬的所有調用，您可以過濾與您相關的調用。