[Pytest 101] 03 Mock SQL Query
Free your unit tests from database dependencies using mocks.
With mock objects, we can simulate database operations without actually connecting to a real database.
This means we don’t need to prepare a dedicated test database or worry about cleaning it up after tests.
Defining a Mock Object Class
Before we start, let’s talk about how to define mock object classes.
By default, Mock or MagicMock returns a new mock object for any accessed attribute or method, even if it doesn’t exist.
While convenient, this behavior can sometimes be too permissive and unrealistic.
To avoid this, you can use the spec
argument to enforce a specification on the mock.
It ensures that the mock mimics only the defined attributes/methods of the original class.
Trying to access an undefined attribute will raise an AttributeError, making the mock stricter and more accurate for unit testing.
First, import the class you want to mock:
from model.user_table import User
Then, use spec when creating the mock:
mock_user = mocker.MagicMock(spec=User)
Mocking Chained Calls
The following feature simplifies mocking for methods that return other mocks:
Each time you access return_value before calling a method, a new mock is created.
Suppose you have a query like this:
user = session.query(User).filter(User.id == user_id).first()
Here’s how to mock each step manually:
# Mock the User object
mock_user = mocker.MagicMock(spec=User)
# Mock the session
mock_session = mocker.MagicMock()
# session.query(...) returns mock_query
mock_query = mocker.MagicMock()
mock_session.query.return_value = mock_query
# query.filter(...) returns mock_filter
mock_filter = mocker.MagicMock()
mock_query.filter.return_value = mock_filter
# filter.first() returns mock_user
mock_filter.first.return_value = mock_user
# Call the function under test
user = get_user_by_id(mock_session, user_id=1)
In this example, we use MagicMock to simulate:
mock_session: The database session object.mock_query: The query object.mock_filter: The return value of thefiltermethod of the query object.mock_user: A mock user instance.
If you’re not testing the individual query steps, you can simplify it using chained mocks:
mock_user = mocker.MagicMock(spec=User)
mock_session = mocker.MagicMock()
mock_session.query.return_value.filter.return_value.first.return_value = mock_user
# Call the function under test
user = get_user_by_id(mock_session, user_id=1)
Querying Data
Use mocks to simulate database query results and avoid actual connections that could affect test results.
# my_project/service/user_service.py
def get_user_by_id(session, user_id):
user = session.query(User).filter(User.id == user_id).first()
return user
If the user exists, we can assert that the returned result is the same as the mocked object:
# my_project/test/service/test_user_service.py
import pytest
from pytest_mock import MockFixture
from model.user_table import User
from service.user_service import (
get_user_by_id,
)
def test_get_user_by_id_success(mocker: MockFixture):
mock_user = mocker.MagicMock(spec=User)
mock_user.id = 1
mock_session = mocker.MagicMock()
mock_session.query.return_value.filter.return_value.first.return_value = mock_user
user = get_user_by_id(mock_session, user_id=1)
assert user == mock_user
Data Manipulation
When using SQLAlchemy to add, delete, or update records, you usually call session.commit() to apply changes.
In tests, you can check whether commit() was called and how many times, or even set its return_value if needed.
You can also check whether session.add() or session.delete() were called correctly.
Here’s an example using delete:
# my_project/service/user_service.py
def delete_user(session, user_id):
user = session.query(User).filter(User.id == user_id).first()
session.delete(user)
session.commit()
return user.id
Assertion tips:
called: whether the method was calledcall_count: how many times the method was called
# my_project/test/service/test_user_service.py
import pytest
from pytest_mock import MockFixture
from model.user_table import User
from service.user_service import (
delete_user,
)
def test_delete_user_success(mocker: MockFixture):
mock_user = mocker.MagicMock(spec=User)
mock_user.id = 1
mock_session = mocker.MagicMock()
mock_session.query.return_value.filter.return_value.first.return_value = mock_user
mock_session.delete.return_value = None
user_id = delete_user(mock_session, mock_user.id)
assert mock_session.delete.called
assert mock_session.delete.call_count == 1
assert mock_session.commit.called
assert user_id == mock_user.id
Simulating Database Errors
Since databases can fail, it’s important to simulate errors during testing.
For example, if you use SQLAlchemy’s one_or_none(), but multiple records are returned, it raises a MultipleResultsFound exception.
In the following example, we don’t handle the error in code. We simply return the result.
This helps demonstrate how the test can simulate and validate the exception behavior.
# my_project/service/user_service.py
def get_user_by_id(session, id) -> User:
return session.query(User).filter(User.id == id).one_or_none()
How to test:
- Import
MultipleResultsFoundfromsqlalchemy.orm.exc - Mock the
.one_or_none()to return the exception. - Assert that the result is of type
MultipleResultsFound.
# my_project/test/service/test_user_service.py
import pytest
from pytest_mock import MockFixture
from sqlalchemy.orm.exc import MultipleResultsFound
from model.user_table import User
from service.user_service import (
get_user_by_id,
)
def test_get_user_by_id_multiple(mocker: MockFixture):
mock_session = mocker.MagicMock()
mock_session.query.return_value.filter.return_value.one_or_none.return_value = (
MultipleResultsFound()
)
user = get_service_info_by_id(mock_session, 1)
assert isinstance(user, MultipleResultsFound)
