Friday, May 17, 2013

Unit testing your Flask REST Application

If you are already convinced that unit testing is the way to go, skip the next section with the different examples.

Unit Testing Is Not Optional

As it says on the Flask website, "something that is untested is broken." While this obviously isn't always true, there is no way for you to know whether or not you have bugs in your code unless you test it. When you initially start writing your application, you probably think one of the following:
  1. Even if it takes only 10 minutes to write the unit test, that time could be spent writing the next great feature.
  2. Why bother wasting time writing test cases when one can just test the application by visiting the site and clicking on the various pages. Even with a REST API that doesn't have nice user interface, one could theoretically could use something like the Advanced REST Client Google Chrome plugin to inspect the JSON that is produced.
Although this is true when you first start out (i.e. the first week or two), you will quickly realize how unmanageable it becomes. How can you be 100% sure that when you change some line of code that it doesn't have a (negative) affect on some other, seemingly unrelated, bit of code? Enter unit testing.

Unit testing is the notion of testing the parts (units) of your code to ensure that they behave as expected. For example, if you have some code that verifies whether a password is correct for a given user, then the unit test should test all possible scenarios that this function may encounter. For example, what happens when you pass None to your function? What happens when you provide an integer password instead of a string password that you were expecting? What happens when you provide the correct password but for a different user in the system?

By testing each of your components independently, you can prove to yourself (and others!) that everything is behaving as expected. So, if in the future you then change one of the components (say you now have a faster way of searching for users in your database given their first name), you merely have to test it with the code that you have already written to verify that it is doing what is expected. If it returns the same results as before (assuming the new method doesn't change the ordering, of course!) then you know that even though the underlying search algorithm changed, it's still behaving the same way as before. Since it's still behaving the same way as before, any other bit of code that required the searching of users by email will also work just like before. Now you can feel free to pat yourself on the back and sleep easy that night.

The key to good unit testing is to ensure that you truly test each unit independently. In other words, one test should not have an affect on another. For example, if you have a test for adding a new user to a database, that test should not affect the data being used for another test. Thus, it's imperative that you setup and initialize dependencies (like databases) in each of the tests so that you know exactly what the starting point of the test is. If this wasn't the case, you will run into problems when you try to run, for instance, a test for counting the number of users in your database.

Although, the examples that I am providing below are in the context of Flask application development, the principles of unit testing are framework and language independent.

Now that I hopefully convinced you that unit testing is not something that should be an afterthought, but something that should be part of your day-to-day programming habits, let's take a look at some unit testing for Flask application.

Unit Testing Examples

To get started, I would highly recommend the great introduction to the basics of unit testing Flask applications on the Flask website. The examples that I have linked below are for slightly more complicated or specific things that I had to test.

Unit Testing Flask File Uploads Without Any Files

Uploading files is one of those things that pretty much all websites support. Whether it is to merely upload a profile picture, or something more complex like uploading the result of a biological experiment, the startpoint is the same -- the file is on the client's computer and needs to end up on your server.

Since uploading files can be such a crucial part or a web application, it should be tested like any other part of the system. Unfortunately, performing this test without actually uploading a real file (i.e. simulating the entire thing) is something that isn't as straightforward as I initially expected (now that I know how to do it, it's really easy!).

In any case, the code below simulates uploading a file using StringIO and then simulates the FileStorage used by Flask (and Werkzeug) by returning a "mocked" TestingFileStorage of our choosing.
from StringIO import StringIO
import unittest

from flask import Request
from werkzeug import FileStorage
from werkzeug.datastructures import MultiDict

# Import your Flask app from your module
from myapp import app

class FlaskAppUploadFileTestCase(unittest.TestCase):

    def setUp(self):

        app.config['TESTING'] = True
        app.config['CSRF_ENABLED'] = False
        self.app = app

        # .. setup any other stuff ..

    def runTest(self):

        # Loop over some files and the status codes that we are expecting
        for filename, status_code in \
                (('foo.png', 201), ('foo.pdf', 201), ('foo.doc', 201),
                 ('foo.py', 400), ('foo', 400)):

            # The reason why we are defining it in here and not outside
            # this method is that we are setting the filename of the
            # TestingFileStorage to be the one in the for loop. This way
            # we can ensure that the filename that we are "uploading"
            # is the same as the one being used by the application
            class TestingRequest(Request):
                """A testing request to use that will return a
                TestingFileStorage to test the uploading."""
                @property
                def files(self):
                    d = MultiDict()
                    d['file'] = TestingFileStorage(filename=filename)
                    return d

            self.app.request_class = TestingRequest
            test_client = self.app.test_client()

            rv = test_client.post(
                '/files',
                data=dict(
                    file=(StringIO('Foo bar baz'), filename),
                ))
            self.assertEqual(rv.status_code, status_code)
Let's take a look at this code in a bit more detail. The first thing we do in our runTest() method is loop over 5 different file types. The assumption is that our application accepts the first 3, while rejecting the last 2. In the loop, we create a class TestingRequest that we will use as our request class for our application. What this does, is it overrides the files attribute to return a TestingFileStorage (defined below) instead of the FileStorage that is normally returned. As I mentioned in the comments, we are creating this class inside the for loop because we need to set the filename that is returned by the TestingFileStorage equal to the one that we are currently using in the loop.

Now that we have created our custom Request, we tell the Flask app to use ours instead and then create a TestClient. Note, you must set the request_class of the app before you create the TestClient. Using the patched TestClient, we can the POST a "file" as normal. Except instead of using a real file, we use a StringIO object so that we don't actually have to have any random files in our project for testing.

That's it really, using the above code (and the TestingFileStorage below) you should be able to test your file uploading routes without actually having to have any files on disk!

I left the implementation of the TestingFileStorage until the end because I copied and pasted it from the Flask-Uploads extension. So that you don't have to got digging around the source code there, I've copied it here for your reference. Enjoy.
class TestingFileStorage(FileStorage):
    """
    This is a helper for testing upload behavior in your application. You
    can manually create it, and its save method is overloaded to set `saved`
    to the name of the file it was saved to. All of these parameters are
    optional, so only bother setting the ones relevant to your application.

    This was copied from Flask-Uploads.

    :param stream: A stream. The default is an empty stream.
    :param filename: The filename uploaded from the client. The default is the
                     stream's name.
    :param name: The name of the form field it was loaded from. The default is
                 ``None``.
    :param content_type: The content type it was uploaded as. The default is
                         ``application/octet-stream``.
    :param content_length: How long it is. The default is -1.
    :param headers: Multipart headers as a `werkzeug.Headers`. The default is
                    ``None``.
    """
    def __init__(self, stream=None, filename=None, name=None,
                 content_type='application/octet-stream', content_length=-1,
                 headers=None):
        FileStorage.__init__(
            self, stream, filename, name=name,
            content_type=content_type, content_length=content_length,
            headers=None)
        self.saved = None

    def save(self, dst, buffer_size=16384):
        """
        This marks the file as saved by setting the `saved` attribute to the
        name of the file it was saved to.

        :param dst: The file to save to.
        :param buffer_size: Ignored.
        """
        if isinstance(dst, basestring):
            self.saved = dst
        else:
            self.saved = dst.name

REST App Response Status Code Testing Harness

When creating a web service that has different types of users with different privileges, it is useful to be able to quickly test whether the route that you just created returns the correct response code (e.g. 200, 401, 403). While you can obviously copy and paste a lot of code to make this work, I ended up creating a base class that does the job of initializing the database, getting the various kinds of users, and then running all of the tests at once. In essence, I wanted to be able to write a TestCase that looks something like this:
class SomeRouteTestCase(FlaskAppRouteStatusCodeTestCase):
    """Test for /foo"""

    __GET_STATUS_CODES__ = dict(
        user=200,
        admin=200,
        super_user=200
    )

    __PUT_STATUS_CODES__ = dict(
        user=403,
        admin=200,
        super_user=200
    )

    def get(self, user, test_data, db_data):
        self.login(user.email, user.password)
        rv = self.app.get('/foo')
        self.logout()
        return rv

    def post(self, user, test_data, db_data):
        self.login(user.email, user.password)
        rv = self.app.post(
            '/foo',
            data=json.dumps(dict(bar="barbar", bam="bambam")),
            content_type='application/json')
        self.logout()

    @unittest.skip("No PATCH")
    def test_patch(self):
        """Override the PATCH tester since /foo can't be patched."""
        pass
What kind of magic is this? Well, not any kind, really. A TestCase for a particular route simply defines the correct status codes for the various HTTP methods and the code to make the calls. In the example above, anyone should be able to perform GET /foo (all of them have a response code of 200), while only admins and super_users are allowed to POST. Since this route doesn't accept the PATCH method, we are telling unittest to skip it with the @unittest.skip decorator. Using this framework, one can test any route with all types of users with minimal effort.

In order to make this all work, we have to define the FlaskAppRouteStatusCodeTestCase that this TestCase inherits from. Fortunately for you, the basic structure of it is in the gist below and pretty straight forward. If you simply fill out the methods that initialize the database and get the set of users to perform the testing on, you can create tests for your routs with reckless abandon.

Although I have only added the methods in FlaskAppRouteStatusCodeTestCase for GET, POST, and PATCH, it should be trivial to add in any other methods for things like DELETE, PUT, etc. For your viewing pleasure, I've included the full example here: