使用pytest-mock模拟python中的数据库调用

Question

我定义了这个函数：

def count_occurrences(cursor, cust_id):
    cursor.execute("SELECT count(id) FROM users WHERE customer_id = %s", (cust_id))
    total_count = cursor.fetchone()[0]
    if total_count == 0:
        return True
    else:
        return False

我想要对它进行单元测试，因此我需要在这里模拟数据库调用。

如何使用pytest，mock来做到这一点？

Answer 1

因此，测试相对简单，因为您的函数接受一个cursor对象，我们可以将其替换为Mock对象。然后，我们所要做的就是configure我们的模拟游标，以返回不同的数据来测试我们的不同场景。

在您的测试中，有两种可能的结果，True或False。因此，我们使用pytest.mark.parametrize为fetchone提供不同的返回值来测试这两种结果，如下所示。

@pytest.mark.parametrize("count,expected", 
    [(0, True), (1, False), (25, False), (None, False)]
)
def test_count_occurences(count, expected):
    mock_cursor = MagicMock()
    mock_cursor.configure_mock(
        **{
            "fetchone.return_value": [count]
        }
    )

    actual = count_occurrences(mock_cursor, "some_id")
    assert actual == expected

当我们运行它时，我们看到针对所有提供的输入运行了四个单独的测试。

collected 4 items                                                                                                

test_foo.py::test_count_occurences[0-True] PASSED                                                          [ 25%]
test_foo.py::test_count_occurences[1-False] PASSED                                                         [ 50%]
test_foo.py::test_count_occurences[25-False] PASSED                                                        [ 75%]
test_foo.py::test_count_occurences[None-False] PASSED                                                      [100%]

=============================================== 4 passed in 0.07s ================================================