使用 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)

这些是对模拟的所有调用，您可以过滤与您相关的调用。