DEV Community

David Sánchez
David Sánchez

Posted on • Edited on

Mocking and patching in Python for unit testing

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
    # ...
Enter fullscreen mode Exit fullscreen mode

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)
Enter fullscreen mode Exit fullscreen mode

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
Enter fullscreen mode Exit fullscreen mode

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()
Enter fullscreen mode Exit fullscreen mode

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
    # ...
Enter fullscreen mode Exit fullscreen mode

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)
Enter fullscreen mode Exit fullscreen mode

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 ❤️.

Top comments (1)

Collapse
 
lawrencek1992 profile image
Kelly

This is great man! We are using django_s3_storage to do a variety of things, including storing a number of files on Cloudfront. We honestly should think about writing some integration tests, and moto imo would be the move for that. But I really just wanted an example how to mock the django s3 stuff in order to test a method that stores an image and then does a few other things in our db. This is perfect. You saved me a ton of time thinking through this on my own. Appreciate it!