Mocking SMTP with pytest-mock
tl;dr
Use the pytest-mock plugin to patch
smtplib.SMTP
# add 'mocker' fixture to your test arguments def test_send_email(mocker): """Function to test your email sending function""" # mock the smtplib.SMTP object mock_SMTP = mocker.MagicMock(name="your_module.smtplib.SMTP") mocker.patch("your_module.smtplib.SMTP", new=mock_SMTP) # run your send email function send_email() # make some assertions. e.g. assert mock_SMTP.return_value.send_message.call_count == 1
Full code explanation further down the article đ
The Problem
Writing unit-tests is an important part of the development process. It helps build trust with end users and colleagues, as well as helping you write better code in the first place .
As part of your data pipeline you might want to send an email to stakeholders. To either notify them that the data is ready for consumption or if there has been any issues processing the data.
How do you write a test for your code which sends the email?
Writing unit-tests for code which interacts with external dependencies (such as an email server), can be tricky.
You could try writing a test which connects to your email server and sends real email, like you would in production. But this comes with a few problems:
- Spamming inboxes with test emails. Every time you run the test, you are actually sending an email to someone. Filling the inboxes of your stakeholders (or even developer accounts) is not ideal.
- Your test isn’t isolated. It relies on an active connection to the email server. If the external email server goes down, your test will start failing unexpectedly even though the issue is nothing to do with the code.
- You can’t reliably emulate exceptions. If your code has error handling in place (for when the email server goes down), you cannot reliably reproduce a specific error when connected to a live server. Therefore you cannot write tests for your exception handling code, reducing the code test coverage.
To decouple your tests from external dependencies we can utilise ‘mocking’.
Mocking
Mocking involves substituting a real object for a ‘fake’ object which behaves like the original. This allows us to emulate the original behaviour, without needing to actually call the real object with an external dependency.
In the example of sending emails, we can ‘mock’ the object which connects to the external email server (smptlib.SMTP
) and any calls we make on that object. For example creating the connection to the server and then sending the email.
Once we have ‘mocked’ the email server object, the code will behave as though it has sent an email but just won’t actually send it.
We can also trigger errors and exceptions (side effects) from the mocked object. This can be very useful for recreating various scenarios and allow you to test your exception handling code.
pytest-mock
We can ‘mock’ the email server using the pytest-mock plugin.
pytest-mock
is a thin wrapper for the patching API provided by Python’s mock package
.
There isn’t much documentation on pytest-mock
itself, but as it is a wrapper for Python’s mock, we can utilise the unittest.mock
documentation to understand the functionality.
How to use pytest-mock to mock sending emails in your unit tests
The best way to explain is to run through an example.
đ» All code examples are available in the e4ds-snippets GitHub repo
Note, the focus of this post in on unit-testing the code that sends an email, rather than how to send an email itself.
I might write another post later on how to send emails, but for now, let’s use the script below as an example for us to unit-test.
The function below that we specifically want to test is the send_pipeline_notification
function.
Example Script
# send_email.py
"""Example code for sending an email"""
import smtplib
import sys
from configparser import ConfigParser
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
def run_pipeline():
"""Some dummy code to run data pipeline"""
# pipeline code ....
# example statistics generated from the run for email reporting purposes
summary = {
"total_files": 100,
"success": 100,
"failed": 0,
"output_location": "gcs://my_data_lake/processed/example.csv",
}
return summary
def build_success_email_notification(
config: ConfigParser, summary: dict[str, int]
) -> MIMEMultipart:
"""Build MIME message object"""
sender_email = config.get("email", "sender_email")
receiver_email = config.get("email", "receiver_email")
msg = MIMEMultipart()
msg["From"] = sender_email
msg["To"] = receiver_email
msg["Subject"] = "Completion Notification"
email_body = f"""
The pipeline has completed successfully.
Total files processed: {summary.get('total_files')}
Successful: {summary.get('success')}
Failed: {summary.get('failed')}
Output location: {summary.get('output_location')}
"""
msg_body = MIMEText(email_body, "plain")
msg.attach(msg_body)
return msg
def send_pipeline_notification(config: ConfigParser, msg: MIMEMultipart) -> None:
"""Send an email, include some error handling"""
host = config.get("email", "host")
port = int(config.get("email", "port"))
try:
with smtplib.SMTP(host, port) as server:
server.send_message(msg)
print("Email notification sent successfully")
except smtplib.SMTPException:
print("SMTP Exception. Email failed to send")
def main():
# read configuration file parameters
config = ConfigParser()
config.read("pipeline_config.ini")
# print config to console for reference
config.write(sys.stdout)
# run pipeline code and return the summary info to report in the email
summary = run_pipeline()
msg = build_success_email_notification(config, summary)
send_pipeline_notification(config, msg)
if __name__ == "__main__":
main()
# pipeline_config.ini
[email]
sender_email = person1@company.com
receiver_email = person2@company.com
host = localhost
port = 1025
The full script works as follows:
- Read some config parameters from a file (e.g. email server settings, sender and receiver email addresses)
- Runs a ‘pipeline’ with some sudo code to imitate a real pipeline
- Sends a notification email with some summary statistics about the run once the pipeline has finished
Let’s also look at the send_pipeline_notification
function that we want to test in more detail:
- Try connecting to an email server using smtplib.SMTP
- Send the email message using the send_message method
- After sending the email we print to the console a message to say the email has sent successfully
- There is also some exception handling in place in case there is an issue sending the email. In that scenario, we simply print that the email failed to send
Running the example script and sending an email using Python
For this demo we will just set up a local SMTP server (localhost, port=1025) to demonstrate sending the email.
In production you would connect to a real email server such as your Gmail account . But the focus of this post is unit-testing rather than setting up connections to real email servers. Therefore, we will keep this part as simple as possible.
To run a SMTP server locally, open up a terminal and run the following command. This will start the server:
python -m smtpd -c DebuggingServer -n localhost:1025
Now open a different terminal and run the example script :
python send_email.py
After running the script, you will notice the email notification will be printed to the console where you are running the SMTP server. This demonstrates that an email has been sent successfully.
Writing a unit-test to mock smptlib.SMTP
Mocking the SMTP object
We want to mock the smtplib.SMTP
object which is used to connect to the email server and send the email.
First ensure you have installed the pytest-mock
dependency into your environment:
pip install pytest-mock
Once installed, you will have access to the mocker
fixture
which we add to our test function arguments.
We can use the mocker
fixture to mock (also known as ‘patching’) the SMTP object as follows:
# test_send_email.py
from send_email import send_pipeline_notification
# add 'mocker' fixture to your test arguments
def test_send_email(mocker):
# mock the smtplib.SMTP object
mock_SMTP = mocker.MagicMock(name="send_email.smtplib.SMTP")
mocker.patch("send_email.smtplib.SMTP", new=mock_SMTP)
# run the send email function
send_pipeline_notification(...)
# make some assertions
assert ....
Let’s go through the ‘patching’ code in more detail:
- Add the ‘mocker’ fixture to our test function arguments
- Create our mocking object using MagicMock
- Patch
smtplib.SMTP
to override the original functionality.
We pass the name of the object we want to patch as a string to the mocker.patch
method. You’ll also notice we prefixed the module name where the object came from: send_email.smtplib.SMTP
.
This prefix is necessary to ensure you are mocking the correct object, avoiding unexpected mocking behaviour .
After patching, whenever smtplib.SMTP
is called during our testing function, the ‘mocked’ object will now be used instead of the real one.
When we call the send_pipeline_notification
function in our test, the code will pretend like it has genuinely sent an email but in fact it has just mocked the functionality.
If you have used mocking before (e.g. for mocking API calls), you might have set a ‘return_value’ for your mocked object. This is used to replicate the expected data returned by the API so it can be used later on in your code. However, when sending emails using SMTP we don’t expect any data back so there is no need to set a return value.
Mocking Exceptions
As mentioned previously, mocking is useful for replicating different exceptions allowing you to test your exception handling.
In our example script we have exception handling when a smtplib.SMTPException
is encountered.
We can test this part of the code by deliberately triggering an exception. This is achieved by setting a ‘side_effect’:
# mock the smtplib object
mock_SMTP = mocker.MagicMock(name="send_email.smtplib.SMTP")
mocker.patch("send_email.smtplib.SMTP", new=mock_SMTP)
# specify what exception should be triggered
mock_SMTP.side_effect = smtplib.SMTPException
After adding the ‘side_effect’, when we call smtplib.SMTP
in the code it will raise the SMTPException
and then run through the exception handling code.
Making assertions
When writing tests, it is not good enough just to test the code runs through without errors. We also want to assert various things about the output to make sure the code is doing what we think it should be doing.
There are various assertion methods built into the mocker
fixture that we can utilise. These are available in the unittest.mock documentation
.
Some useful assertions specific to the mocker
fixture include:
- assert_called() : Asserts that the mocked object was called at least once
- call_count : Records the number of times a method was called
For example, in our send_pipeline_notification
function we want to assert that the send_message
method was called once within the context manager. We would do the following:
assert mock_SMTP.return_value.__enter__.return_value.send_message.call_count == 1
This might look quite convoluted, and it caught me out when I was first trying to test the code. But the reason is because we are using a context manager to instantiate our SMTP object.
When you use a context manager in Python
you first call the __enter__
method, before calling further methods.
We can observe this behaviour if we print all the calls the mocked object made which is stored within mock_SMTP.mock_calls :
# output from print(mock_SMTP.mock_calls)
[call('localhost', 1025),
call().__enter__(),
call().__enter__().send_message(....),
call().__exit__(None, None, None)]
Therefore, our mocked object first calls __enter__
and then it is the returned object from that which calls send_message
. So we have to work our way down the chain to get to the method we want to actually assert.
If you are not using a context manager to instantiate your SMTP object you can simply use:
assert mock_SMTP.return_value.send_message.call_count == 1
We can also use generic Pytest assertions to test any other code in our function.
For example, in our example script we have some print statements which trigger after the email is sent. We can test this using the capsys fixture which records the information printed to the terminal.
# add the capsys fixture to your function arguments
def test_send_pipeline_notification(mocker, capsys):
# setup code ....
output = capsys.readouterr()
assert "Email notification sent successfully" in output.out
Putting it together
Below are the full unit tests we can use to test our email sending code.
We write two tests. One to test the ‘happy path’ where the email sends successfully. And one to test the exception handling code when we receive a smtplib.SMTPException
.
# test_send_email.py
"""Test your email code"""
import smtplib
from configparser import ConfigParser
from send_email import build_success_email_notification, send_pipeline_notification
def test_send_pipeline_notification(mocker, capsys):
"""Test the successful sending of the email"""
# 1. SETUP
# 1a. read config
config = ConfigParser()
config.read("pipeline_config.ini")
# 1b. build an example message
summary = {
"total_files": 1_000,
"success": 1_000,
"failed": 0,
"output_location": "gcs://my_data_lake/processed/",
}
msg = build_success_email_notification(config, summary)
# 1c. mock the smptlib.SMTP object
mock_SMTP = mocker.MagicMock(name="send_email.smtplib.SMTP")
mocker.patch("send_email.smtplib.SMTP", new=mock_SMTP)
# 2. ACT -- Run the function
send_pipeline_notification(config, msg)
# 3. ASSERT -- Test the outputs
# 3a. Test 'send_message' method was actually called
# Use the below if you are using a context manager
assert mock_SMTP.return_value.__enter__.return_value.send_message.call_count == 1
# Use the below if you are NOT using a context manager
# assert mock_SMTP.return_value.send_message.call_count == 1
# 3b. Test print statement after successfully running 'send_message'
output = capsys.readouterr()
assert "Email notification sent successfully" in output.out
def test_send_pipeline_notification_smtpexception(mocker, capsys):
"""Test SMTPException exception handling"""
# 1. SETUP
# 1a. read config
config = ConfigParser()
config.read("pipeline_config.ini")
# 1b. build an example message
summary = {
"total_files": 1_000,
"success": 1_000,
"failed": 0,
"output_location": "gcs://my_data_lake/processed/",
}
msg = build_success_email_notification(config, summary)
# 1c. mock the smptlib.SMTP object
mock_SMTP = mocker.MagicMock(name="send_email.smtplib.SMTP")
mocker.patch("send_email.smtplib.SMTP", new=mock_SMTP)
# mock STMPException error
mock_SMTP.side_effect = smtplib.SMTPException
# 2. ACT -- Run the function
send_pipeline_notification(config, msg)
# 3. ASSERT -- Test the outputs
# 3a. Test 'send_message' method was NOT called
# use if the below you are using a context manager
assert mock_SMTP.return_value.__enter__.return_value.send_message.call_count == 0
# use the below if you are NOT using a context manager
# assert mock_SMTP.return_value.send_message.call_count == 0
# 3b. Test print statement after triggering SMTPException
output = capsys.readouterr()
assert "SMTP Exception. Email failed to send" in output.out
You can run these tests by simply running:
pytest test_send_email.py
If you still have your local SMTP sever running, you will notice that when you run your unit-tests, the output in the terminal will not update.
This demonstrates that the mocking has overridden the SMTP object and a real email is no longer being sent.
If you were to comment out the line of code where mocker.patch
is called, and re-ran the test cases, then the SMTP server terminal output would update. This shows you are back to sending real emails again.
Conclusion
In this post we explained how to mock the smtplib.SMTP
object in your Pytest unit tests:
- We created a simple script to run a ‘pipeline’ and send an email at the end of it.
- We set up a simple SMTP server locally to send this email and demonstrate a real connection to an email server in production.
- We used
pytest-mock
to patch the functionality ofsmtplib.SMTP
and demonstrated how to manipulate the object’s behaviour such as forcing specific Exceptions. - We wrote some unit tests to test the
send_pipeline_notification
function which made assertions on the expected outputs.
I hope this post was useful.
When I first tried to mock smtplib.SMTP
I couldn’t find much documentation on how to do it in Pytest with pytest-mock
. All of the mocking documentation related to using the ‘unittest’ library instead.
I found it confusing at first but now I realise pytest-mock
and unittest.mock
are basically the same thing. Any documentation and usage is pretty consistent between the two.
I also got caught out when trying to make assertions for objects with a context manager. I was scratching my head for ages working out why assert mock_SMTP.return_value.send_message.call_count
was returning 0 even though, the method was clearly being called đ€·â.
I hope documenting this process saves you some time and helps increase your test coverage for code with external dependencies.
Happy coding!
đ» All code examples are available in the e4ds-snippets GitHub repo
Resources
- StackOverflow discussion
- Toptal - Introduction to mocking in Python
- RealPython - Sending emails using Python
- smtplib documentation
- unittest.mock documentation - useful for finding attributes returned by the test mocking object
Further Reading
- Top tips for using PyTest
- Unit testing PySpark code using Pytest
- How to save the output of PySpark DataFrame 'show' to a variable
- Google Search Console API with Python
- Deploying Dremio on Google Cloud
- Gitmoji: Add Emojis to Your Git Commit Messages!
- Five Tips to Elevate the Readability of your Python Code
- Do Programmers Need to be able to Type Fast?
- How to Manage Multiple Git Accounts on the Same Machine