Saturday, January 26, 2013

Multi-Threaded Unit Test for (Flask) REST APIs

I've been working on creating a RESTful API using the Flask microframework recently, and had to come up with a way to test it using multiple concurrent threads (i.e. the way it will actually be used when it goes live). Although this example does not preclude the use of the test client that should be used for unit testing the individual parts of your site, this method can be used in addition to your other tests to make sure that it can handle multiple simultaneous requests. Also, as with most unit tests, this is a simulation and does not exactly mirror all of the details of an actual deployment. In particular, this test makes use of the built in server, which is not intended to be a production server.

Apart from the imports that you will need to make your application work, we will be making use of the fabulous Requests package to handle all of the interaction with the web service from within our unit test. While one can do all of this using urllib and urllib2, Requests makes our lives quite a lot easier. For example, here I created some helper methods that allow me to log in and log out of the web service. (For those of you who think that this isn't RESTful since we have to keep some state on the server to allow for logging in/out, I find that this small deviation from a true RESTful interface is justified as it makes everything else much simpler and cleaner.)
def login(self, username, password):
    """Helper method to let us log in to our web service."""

    # Create a dictionary of login data
    login_data = json.dumps(dict(username=username, password=password))
 
    # Log in to our service
    return requests.post(SERVER_URL + "/login", data=login_data,
                                headers={'content-type': 'application/json'})
 
def logout(self, cookies):
    """Helper method to let us log out from our web service."""
    return requests.post(SERVER_URL + "/logout", cookies=cookies)
The real "magic" happens in the following code. We're simply setting up a thread that starts the Flask application. Note, you have to set threaded to True, to ensure that the built in web server code doesn't just run in a single thread.
def start_and_init_server(app):
    """A helper function to start out server in a thread.

    This could be done as a lamnda function, but this way we can
    perform other setup functions if necessary.

    Args:
        app: The Flask app to run
    """
    app.run(threaded=True)
 
# Create a thread that will contain our running server
server_thread = Thread(target=start_and_init_server, args=(self.app, ))
server_thread.start()
Now that the server has been started in its own thread, you can simply bombard it with as many requests as you'd like. For example, if we are to assume that this REST interface is for a blogging system that allows for me to POST a new blog post to /posts, then we can do something like the following:
for i in range(n_new_posts):
    t = Thread(target=post_data)
    t.start()
Of course, the above example assumes that you have some function called post_data that does the job of POSTing some data to the web service.

Again, this isn't a foolproof method to test your code as there may be some issues that come up with you deploy it on your development server, but this should be a good start. And hey, you can always write yourself a test that fires up that actual server that you will be using instead of the one that is included with Flask (or your web framework of choice).

Now that you've seen the bits and pieces, you can find the full example as a gist on github. As aforementioned, we're assuming that we are creating a simple blogging API and that there are a couple REST endpoints such as /login, /logout, /posts, etc. You'll obviously want to change those endpoints to meet your needs.

In case you are too lazy to head over to github, I've also included the code here for your convenience.

No comments:

Post a Comment