[Pytest 101] 04 Mock Class
An introduction to two approaches for testing class instantiation when using the Factory Pattern, with practical examples using pytest-mock.
Introducing the factory pattern in your code can significantly improve its structure—but it also makes unit testing a bit trickier.
We’ll walk through two different mocking strategies, compare their pros and cons, and help you write more robust tests!
When using the Factory Design Pattern to return different classes based on logic, there are two typical ways to write unit tests.
# my_project/service/alert_service.py
from service.alert.mail_alert import MailAlert
from service.alert.teams_alert import TeamsAlert
from service.alert.discord_alert import DiscordAlert
def get_alert_model(session, alert_type, alert_config):
if alert_type == "mail":
return MailAlert(session, alert_config)
elif alert_type == "teams":
return TeamsAlert(session, alert_config)
elif alert_type == "discord":
return DiscordAlert(session, alert_config)
else:
raise Exception(
f"Alert Type is not valid. Please check the alert type mail, teams or discord. {alert_type}"
)
Mock class
Since the classes are already imported into the module, you need to mock the versions inside the target module (i.e., alert_service.py):
- service.alert_service.MailAlert
- service.alert_service.TeamsAlert
- service.alert_service.DiscordAlert
# my_project/test/service/test_alert_service.py
import pytest
from pytest_mock import MockFixture
from service.alert_service import (
get_alert_model,
)
def test_get_alert_model_success(mocker: MockFixture):
mock_session = mocker.MagicMock()
alert_config = {}
alert_type_model_dict = {
"mail": "MailAlert",
"teams": "TeamsAlert",
"discord": "DiscordAlert",
}
for alert_type, alert_model in alert_type_model_dict.items():
mock_alert_model = mocker.patch(f"service.alert_service.{alert_model}")
result = get_alert_model(mock_session, alert_type, alert_config)
assert result == mock_alert_model.return_value
mock_alert_model.assert_called_once_with(mock_session, alert_config)
def test_get_alert_model_invalid_type(mocker: MockFixture):
mock_session = mocker.MagicMock()
alert_type = "test_type"
with pytest.raises(
Exception,
match=f"Alert Type is not valid. Please check the alert type mail, teams or discord. {alert_type}",
):
get_alert_model(mock_session, alert_type, {})
Explanation
Using mocker.patch
, we mock the appropriate object into mock_alert_model and then assert that the function returns mock_alert_model.
We also use assert_called_once_with to make sure the mocked class was only called once and with the expected arguments.
For more assertion methods, see: Mock examples
Importing the Actual Classes
Another approach is importing the actual classes into the test, which creates a real dependency on them.
This approach better mimics runtime behavior, but adds complexity to the test itself.
This method allows for using isinstance to check the returned type. However, changes to the real class may break tests, and if the input config is insufficient to construct the real class, it will raise an error.
For example, if class construction depends on alert_config inside alert_content but we pass in an empty dict, it will raise:
AttributeError: ‘dict’ object has no attribute ‘alert_content’
# my_project/test/service/test_alert_service.py
import pytest
from pytest_mock import MockFixture
from service.monitor.alert.mop_mail_alert import MailAlert
from service.monitor.alert.teams_alert import TeamsAlert
from service.monitor.alert.trigger_airflow_dags import DiscordAlert
from service.alert_service import (
get_alert_model,
)
def test_get_alert_model_success(mocker: MockFixture):
mock_session = mocker.MagicMock()
alert_config = {}
# Use actual classes
alert_type_model_dict = {
"mail": MailAlert,
"teams": TeamsAlert,
"discord": DiscordAlert,
}
# No mocking for classes
for alert_type, alert_model in alert_type_model_dict.items():
result = get_alert_model(mock_session, alert_type, alert_config)
assert isinstance(result, alert_model)
