Python Mocking in Production
Unit testing is an art. While the question of what to test is crucial (you can't test everything), in this post we'll take a closer look at some advanced testing and mocking techniques in Python. While some of these were already introduced in previous postings (Part 1, Part 2, Part 3), this post differs by showing their interplay on a real-world example, as well as goes above the scope of previous posts and adds more insights.
Among others, we will discuss:
- patching functions – with the additional requirements of patching async functions and modifying their behaviour
- asserting that functions are called as expected
- writing custom matchers for partially matching call arguments
Recap of Necessary Topics
Before diving into the actual contents of this post, we first give a quick recap of how to use pytest. Afterwards, we do the same for asyncio, since our example further on makes use of this. If you are already familiar with these topics, please feel free to skip to the next section.
pytest
Unit testing should be part of any somewhat professional software product. It helps avoid bugs and thus increases efficiency and quality of the code. While Python natively ships with unittest, pytest is prefered by many developers.
To begin, let's look at one example from a previous post:
import pytest
class MySummationClass:
def sum(self, x, y):
return x + y
@pytest.fixture
def my_summation_class():
return MySummationClass()
def test_sum(my_summation_class):
assert my_summation_class.sum(2, 3) == 5We can run this test via python -m pytest test_filename.py. When doing so, pytest discovers all tests in that file following some conventions (e.g. all functions named test_...) and executes them. In our case, we defined a fixture returning an instance of MySummationClass. Fixtures can be used to e.g. avoid repetitive initialisation and to moduralize code. We then call that instance's sum() method, and check that the result equals the expected one via assert.
Mocking
Often, during testing we encounter functions we either can't or don't want to run – e.g. because they would take too long or have undesired side effects. For this purpose, we can mock them out.
Let's consider an example from the previous post:
import time
from unittest.mock import Mock, patch
def expensive_function() -> int:
time.sleep(10)
return 42
def function_to_be_tested() -> int:
expensive_operation_result = expensive_function()
return expensive_operation_result * 2
@patch("sample_file_name.expensive_function")
def test_function_to_be_tested(mock_expensive_function: Mock) -> None:
mock_expensive_function.return_value = 42
assert function_to_be_tested() == 84We're using a decorator to patch expensive_function with mock_expensive_function, this way replacing the original function's long run time by a function with similar properties, chosen by us.
Lastly, let's briefly recap asyncio: asyncio is a multi-threading library whose primary area of application is I/O bound applications – that is applications, which spend a large portion of their time waiting for inputs or outputs. asyncio actually uses a single thread for this, and leaves it up to the developer to define when coroutines can yield execution and hand over to others.
Let's re-use the motivating example from the previous post:
import asyncio
async def sleepy_function():
print("Before sleeping.")
await asyncio.sleep(1)
print("After sleeping.")
async def main():
await asyncio.gather(*[sleepy_function() for _ in range(3)])
asyncio.run(main())If we were to run sleepy_function conventionally three times in a row, this would take 3s. However, with asyncio this program finishes in 1s: gather schedules the execution of three function calls, and inside sleepy_function the keyword await yields control back to the main loop, which has time to execute other code (here: other instances of sleepy_function) while sleeping for 1s.
Problem Setup
Now, equipped with sufficient prior knowledge, let's dive deeper into the actual contents of this post. In particular, in this section we first define the Programming problem serving as playground for unit testing.
For setting up the project, we used poetry, and also followed other best practises, such as using typing, formatters and linters.
Our example models generating some messages and sending them via some client (e.g. email): in the main file, we first instantiate the client via a factory function, then generate some messages, and lastly send these messages asynchronously using the client.
The project consists of the following files, which you can also find on github:
pyproject.toml
[tool.poetry]
name = "Pytest Example"
version = "0.1.0"
description = "A somewhat larger pytest example"
authors = ["hermanmichaels <[email protected]>"]
[tool.poetry.dependencies]
Python = "3.10"
mypy = "0.910"
pytest = "7.1.2"
pytest-asyncio = "0.21.0"
black = "22.3.0"
flake8 = "4.0.1"
isort = "^5.10.1"message_sending.py
import asyncio
from message_utils import Message, generate_message_client
def generate_messages() -> list[Message]:
return [Message("Message 1"), Message("Message 2")]
async def send_messages() -> None:
message_client = generate_message_client()
messages = generate_messages()
await asyncio.gather(*[message_client.send(message) for message in messages])
def main() -> None:
asyncio.run(send_messages())
if __name__ == "__main__":
main()message_utils.py
from dataclasses import dataclass
@dataclass
class Message:
body: str
class MessageClient:
async def send(self, message: Message) -> None:
print(f"Sending message: {message}")
def generate_message_client() -> MessageClient:
return MessageClient()We can run this program via python message_sending.py, which – as stated above – first instantiates a MessageClient, then generates a list of dummy messages via generate_messages, and eventually sends these using asyncio. In the last step, we create tasks message_client.send(message) for every message, and then run these asynchronously via gather.
Testing
With that, let's come to testing. Here, our goal is to create some scenarios, and to ensure that the correct messages are being send out via the message client. Naturally, our simple demo setting is too simplistic to cover this, but imagine the following: in the real, you're using the client to send out messages to customers. Depending on certain events (e.g. product bought / sold), there will be different messages created. You thus want to simulate these different scenarios (e.g. mock someone buying a product), and check, that the right emails are being generated and sent out.
Sending actual emails during testing is probably not desired though: it would put stress on the email server, and would require certain setup steps, such as entering credentials, etc. Thus, we want to mock out the message client, in particular it's send function. During the test we then simply put some expectations on this function, and verify it was called as expected (e.g. with the right messages). Here, we will not mock generate_messages: while certainly possible (and desired in some unit tests), the idea here is to not touch the message generating logic – while obviously very simplistic here, in a real system the messages would be generated based on certain conditions, which we want to test (one could thus call this more of an integration test, than an isolated unit test).
Test Function was Called Once
For a first try, let's change generate_messages to only create a single message. Then, we expect the send() function to be called once, which we will test here.
This is how the corresponding test looks:
from unittest.mock import AsyncMock, Mock, call, patch
import pytest as pytest
from message_sending import send_messages
from message_utils import Message
@pytest.fixture
def message_client_mock():
message_client_mock = Mock()
message_client_mock_send = AsyncMock()
message_client_mock.send = message_client_mock_send
return message_client_mock
@pytest.mark.asyncio
@patch("message_sending.generate_message_client")
async def test_send_messages(
generate_message_client: Mock, message_client_mock: Mock
):
generate_message_client.return_value = message_client_mock
await send_messages()
message_client_mock.send.assert_called_once()Let's dissect this in more details: test_send_messages is our main test function. We patched the function generate_message_client, in order to not use the real (email) client returned in the original function. Pay attention to "where to patch": generate_message_client is defined in message_utils, but since it is imported via from message_utils import generate_message_client, we have to target message_sending as the patch target.
We're not done yet though, due to asyncio. If we continued without adding more details to the mocked message client, we would get an error similar to the following:
TypeError: An asyncio.Future, a coroutine or an awaitable is required … ValueError(‘The future belongs to a different loop than ‘ ‘the one specified as the loop argument').
The reason for this is that in message_sending we call asyncio.gather on message_client.send. However, as of now, the mocked message client, and consequently its send message, are simply Mock objects, which cannot be scheduled asynchronously. In order to get around this, we introduced the fixture message_client_mock. In this, we define a Mock object called message_client_mock, and then define its send method as an AsyncMock object. Then, we assign this as return_value to the generate_message_client function.
Note that pytest natively actually does not support asyncio, but needs the package pytest-asyncio, which we installed in the pyproject.toml file.
Test Function was Called Once With Specific Argument
As a next step, we not only want to check send was called once, as expected – but also ensure it was called with the right arguments – i.e. the right message.
For this, we first overload the equals operator for Message:
def __eq__(self, other: object) -> bool:
if not isinstance(other, Message):
return NotImplemented
return self.body == other.bodyThen, at the end of the test, we use the following expectation:
message_client_mock.send.assert_called_once_with(Message("Message 1"))Partially Matching Arguments
Sometimes, we might want to do some "fuzzy" matching – that is, don't check for the exact arguments a function was called with, but check some portion of them. To stay with our user story of sending emails: imagine, the actual emails contains lots of text, of which some is somewhat arbitrary and specific (e.g. a date / time).
To do this, we implement a new proxy class, which implements the __eq__ operator w.r.t. Message. Here, we simply subclass string, and check it being contained in message.body:
class MessageMatcher(str):
def __eq__(self, other: object):
if not isinstance(other, Message):
return NotImplemented
return self in other.bodyWe can then assert that the sent message e.g. contains a "1":
message_client_mock.send.assert_called_once_with(MessageMatcher("1"))Checking Arguments for Multiple Calls
Naturally, only being able to check that a function was called once is not very helpful. If we want to check a function was called N times with different arguments, we can use assert_has_calls. This expects a list of of type call, with each element describing one expected call:
message_client_mock.send.assert_has_calls(
[call(Message("Message 1")), call(MessageMatcher("2"))]
)Conclusion
This brings us to the end of this article. After recapping the basics of pytest and asyncio, we dove into a real-world example and analysed some advanced testing, and in particular mocking, techniques.
We saw how to test and mock async functions, how to assert they are called as expected, and how to relax equality checks for the expected arguments.
I hope you enjoyed reading this tutorial – thanks for tuning in!
