loading...

Mocking and patching in Python for unit testing

d4vsanchez profile image David Sánchez Updated on ・5 min read

Ubidots is an effortless point-and-click Internet of Things (IoT) application builder with data analytics and visualization. I am an engineer at Ubidots, and it is our daily goal to seamlessly and cleanly turn your sensor data into information that matters. Hiring an engineering team to develop a platform that both functions and looks great is costly in both time and money so we did it for you. One feature that many clients enjoy is Historical Reporting which is generated and saved with Amazon S3.

When developing this feature, we needed to create a unit test using Python to test feature reliability. Using a mocked class of the boto3 module. As we were in a test environment we couldn't allow the upload of files to our S3 Bucket nor did we have the correct credentials so we needed to find a way to mock the class but including (or actually taking out) the following features:

  1. Don't communicate with Amazon, as I don't need to upload anything.
  2. Don't raise an exception by not having the correct credentials.

For obvious reasons, I cannot copy & paste the Ubidots source code, but for the purpose of this article, I have included an example somewhat similar to that which we use daily.

Let's code!

Let's start with the most important part of the post: The code!

Create the model

This will be our "Report model" and it's got a method that will let us create the report and send it by email.

import StringIO

from django.db import models
from boto3.session import Session

class Report(models.Model):
    # ...
    # Information about the model

    def create_report_and_send_by_email(self):
        stream = StringIO.StringIO()
        self.create_report(stream)

        # Go back to the first bit of the stream
        stream.seek(0)

        session = Session(aws_access_key_id=AWS_ACCESS_KEY, aws_secret_access_key=AWS_SECRET_ACCESS_KEY)
        s3 = session.resource('s3')
        s3.Bucket(AWS_BUCKET_DATA).put_object(
            Key='file.pdf',
            Body=stream.read(),
            ACL='public-read',
            Expires=(datetime.datetime.now() + datetime.timedelta(days=1)).strftime('%Y/%m/%d'),
            ContentType='application/pdf'
        )

        # Do some other stuff and send the file by email

    # More methods
    # ...

Here, I've only implemented one method of our model as this is the focus of this article, but normally this model will have additional fields and other methods which have been left off solely for simplicity reasons.

Create the test

This will be the file that implements the unit tests for our model, you can read a bit more about writing and running unit tests in the Django Documentation site.

from django.core import mail
from django.test import TestCase
from reports.models.report import Report

class ReportTest(TestCase):
    def tearDown(self):
        Report.objects.all().delete()
        mail.outbox = []

    def test_create_report_and_send_by_email(self):
        info = # This is a dict with the information of the fake Report
        report = Report.objects.create(**info)

        # There shouldn't be emails in the outbox
        self.assertEqual(len(mail.outbox), 0)

        report.create_report_and_send_by_email(self)

        # There should be one email in the outbox
        self.assertEqual(len(mail.outbox), 1)

If we were to run this code without the correct AWS credentials, it will raise an exception that says that the credentials are not correct and therefore our test will fail.

Here is when we need to do a mock of the Session class from boto3; preventing the session from making a connection to Amazon S3. In this case, we expect nothing to be returned, the only thing we're going to do with the mock is prevent the connection with Amazon S3.

Create the mock

The first thing we need to do is import the mock and boto3 libraries in our test file.

import mock
from boto3.session import Session

Now, we need to create our mock for Session and Resource classes, and we'll implement the methods we use with the same arguments to run our test. Below you will find our "FakeSession" and "FakeResource" to be used as our mocks classes.

class FakeResource():
    def Bucket(self, bucket):
        return self

    def put_object(self, Key=None, Body=None, ACL='', Expires='', ContentType=''):
        # We do nothing here, but return the same data type without data
        return {}

class FakeSession(Session):
    def resource(self, name):
        return FakeResource()

Now, with these two mock classes created, we can add them to the test case using the patch. Using the patch, the test will run using these fake classes instead of the real ones from boto3.

@mock.patch('reports.models.report.Session', FakeSession)
def test_create_report_and_send_by_email(self):
    # Same implementation we used before
    # ...

As you can see, the patch is made to the namespace of the module that we're testing, not to the namespace of the original module we want to modify.

Adding the mock to the Test

Our test code (including the previous steps) looks like this:

import mock

from django.core import mail
from django.test import TestCase
from boto3.session import Session
from reports.models.report import Report

class FakeResource():
    def Bucket(self, bucket):
        return self

    def put_object(self, Key=None, Body=None, ACL='', Expires='', ContentType=''):
        # We do nothing here, but return the same data type without data
        return {}

class FakeSession(Session):
    def resource(self, name):
        return FakeResource()

@mock.patch('reports.models.report.Session', FakeSession)
class ReportTest(TestCase):
    def tearDown(self):
        Report.objects.all().delete()
        mail.outbox = []

    def test_create_report_and_send_by_email(self):
        info = {"name": "Fake report"} # This is a dict with the information of the fake Report
        report = Report.objects.create(**info)

        # There shouldn't be emails in the outbox
        self.assertEqual(len(mail.outbox), 0)

        report.create_report_and_send_by_email(self)

        # There should be one email in the outbox
        self.assertEqual(len(mail.outbox), 1)

When this new test case is executed it will use the implementations we've created in FakeSession each time Session from boto3 is instanced.


Using the patching decorator we were able to make a mock class from a third-party (boto3) work the way we needed to test some modules; even when missing some parameters that would otherwise be available in a production environment, using a mock class we could test our module without the need of building content data to run the program.

Authors' thoughts: This mock should only be made when it is necessary. It is important to think about how will this modification affect the test of the module; in the case of this piece, we only modified modules to prevent boto3 from connecting to Amazon and did not raise an exception. For our needs, the S3 feature wasn't important for the correct behavior of the test.


If you found this article helpful please leave your reaction and share it in your networks! And you're also very welcome to leave feedback in the comments section ❤️.

Posted on by:

d4vsanchez profile

David Sánchez

@d4vsanchez

Software Engineer. Passionate about technology. Developer by passion. Making cool things at @ubidots. Linux User 🐧. #python #node #react #iot

Discussion

pic
Editor guide