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:

  1. 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.
  2. 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.
  3. 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.

Example email message printed via the local SMTP server
Example email message printed via the local SMTP server

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:

  1. Add the ‘mocker’ fixture to our test function arguments
  2. Create our mocking object using MagicMock
  3. 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 of smtplib.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

Further Reading