Mocking database calls in python using pytest-mock

Question

I have this function defined:

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

I want to unit test it, for that I need to mock the database call here.

How can this be done using pytest, mock?

Answer 1

So this is relatively straightforward to test since your function takes a cursor object which we can substitute with a Mock object. Then all we have to do is configure our mock cursor to return different data to test the different scenarios we have.

In the case of your tests there are two possible outcomes, True or False. So we provide different return values for fetchone to test both of these outcomes as shown below using 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

When we run it we see that four separate tests are ran, against all the inputs provided.

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