python-3.x - 使用 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 来做到这一点？

标签： python-3.xdatabaseunit-testingmockingpytest

Answer 1

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

在您的测试中，有两种可能的结果，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 ================================================