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

因此，测试相对简单，因为您的函数采用游标对象，我们可以用模拟对象替换它。然后我们所要做的就是配置我们的模拟游标以返回不同的数据来测试我们拥有的不同场景。

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

@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 ================================================