A production-grade Machine Learning API using Flask, Gunicorn, Nginx, and Docker — Part 4 (Unit Testing API)

In this final part of the series, we will write unit tests for our Flask API. By this point, we have already explored how to write a simple Flask API, set it up inside docker, and use Flask Blueprints to make it modular.

At the end of this part, our folder structure will look like this:

flask-ml-api
|- api
    |- __init__.py
    |- endpoints
        |- __init__.py
        |- classification.py
    |- models
        |- model.pkl
    |- tests
        |- __init__.py
        |- test_classification.py
    |- app.py
    |- wsgi.py
    |- requirements.txt
    |- Dockerfile
|- nginx
    |- nginx.conf
    |- Dockerfile
|- docker-compose.yml
|- run_unittests.py

Step 1: Writing tests for each endpoint

In our API, we only have one endpoint, which lives in classification.py. However, you may have more endpoints, if you have done so in Part 3.

Let’s start by making a new directory for all our tests. In this directory, I am making a new file for each endpoint. In this case, I have tests/test_classification.py

We will be using the unitest library to write and execute tests.

This is what my test_classification.py looks like:

import unittest
import json
from flask import Flask
from api.endpoints.classification import classification_api

app = Flask(__name__)
app.register_blueprint(classification_api)


class ClassificationTests(unittest.TestCase):

    tester = None

    def __init__(self, *args, **kwargs):
        super(ClassificationTests, self).__init__(*args, **kwargs)
        global tester
        tester = app.test_client()

    def test_classify_single(self):
        response = tester.get(
            '/classification',
            data=json.dumps({"text": "Cocoa setup issues"}),
            content_type='application/json'
        )

        data = response.get_data(as_text=True)
        print("Category predicted: "+str(data))
        self.assertEqual(response.status_code, 200)
        self.assertIsNotNone(data)


if __name__ == '__main__':
    unittest.main()

First of all, the test itself is a Flask app. This is necessary, because our endpoints live as Flask Blueprints, as described in part 3 of this series.

• We call register_blueprint() and pass in the classification_api blueprint we created in classification.py
• Following this, we will store the app.test_client() in a local tester variable. This will give us access to the API, as if we are hitting it with actual traffic.
• Note that in the __init__ method, I had to make sure that tester is a global variable.

Now, you can write the test function, which is test_classify_single in my case. Here, we can use tester.get() to do a GET request on the '/classification' endpoint. The data and content_type are determined by what kind of request your API can handle. In this case, I am passing in a test JSON:

{"text": "Cocoa setup issues"}

This is the JSON that my API should be able to decode, pass into the model which classifies it, and then send back the category as a JSON.

• This returned JSON can be accessed using the response.get_data() method as text that I can then print, or write assert statements on.
• My assert statements are testing to make sure that the status code returned is a success (which should be 200) for this structure of input JSON, and that the data in the response is not None (i.e. null).

IMPORTANT NOTE: With ML APIs, it is bad practice to test the result of the prediction, as updating the model can cause tests to fail even though the API is working perfectly fine. This sort of testing should happen in the validation stage of your model development, rather than after implementing the API. For example, testing whether a statement like “This is a furry feline that purrs” predicts the “cat” class is bad.

Instead, concentrate your testing efforts on making sure that your API can handle weird JSON input, invalid input, and provide correct response codes.

Step 2: Write a script to run all tests for all endpoints

Once we have all our testing scripts ready, we can write a simple script to run all the tests for all endpoints.

My script file is called run_unittests.py and is located in the home folder of the project (alongside docker-compose — see file structure above)

My run_unittests.py looks like this:

import unittest

print("Running unit tests from api/tests directory...")

loader = unittest.TestLoader()
start_dir = 'api/tests'
suite = loader.discover(start_dir)

runner = unittest.TextTestRunner()
runner.run(suite)

print("Running tests is complete")

All this does is create a TestLoader, which loads all test classes from the 'api/tests' folder and runs the entire suite.

This code is adapted from this StackOverflow post.

Step 3: Run the tests!

You can simply use:

python run_unittests.py

to run the tests locally.

If you like, you can also copy the run_unittests.py file into your docker container. You can SSH into your docker container using

docker exec -it <container name> /bin/bash

and then run tests using after to navigate to the directory with this file:

python run_unittests.py

Whoo Hoo! You should see your tests pass (hopefully) or you can go back to your code and check your work.

Conclusion to Part 4

Hope you had fun implementing this tutorial series. You can find all the code I wrote on GitHub.

This is a wrap on the project. There are a lot of online resources that helped me figure this out, and the internet is always your friend if you get stuck. Nonetheless, if you have any questions or feedback, please drop them in the comments and I’ll try to help!

In this series:

Part 1: Setting up our API
Part 2: Integrating Gunicorn, Nginx and Docker
Part 3: Flask Blueprints — managing multiple endpoints
Part 4: Testing your ML API

Hello there! Thanks for reading. Here’s a tad bit about me. I am a Computer Science student at the University of British Columbia, Canada. I primarily work on machine learning projects, mostly NLP. I also do photography as a hobby. You can follow me on Instagram and LinkedIn, or visit my website. Always open to opportunities 🚀