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