Testing with Flask and SQLAlchemy
  |   Source

Testing with PyTest, Flask, and SQLAlchemy

I would like to walk through the process of developing tests using Flask and SQLAlchemy, and shed light on some of the pitfalls and gotchas along the way!

Flask is a popular web framework in Python that is commonly used to build many kinds of web projects. It is pretty easy to hack together everything from small websites to REST API's by leveraging a litany of available extentions (Flask-Restful, Flask-Restplus, Flask-SQLAlchemy, etc.). A very common combination is to use Flask with SQLAlchemy as the ORM. At some point in your career, you are probably going to be required to write tests, whether these are unit tests or integration tests. This may be dictated by project requirements or simply to reduce the complexity of finding bugs in your own code.

Assumptions

I assume you are quite comfortable with:

  • Python
  • Flask
  • App factories
  • SQLAlchemy

And have a working knowledge of:

  • PyTest
  • Unittest

We are not writing an app; we are disecting snippets.

Rabbit holes

  • I'm not here to weigh in on Flask vs. Django vs. [insert your favorite framework]. I am leveraging personal experience discovered on the job. Like me, you may not be able to choose what tools you use at all times.
  • I have a personal preference for PyTest, but you should also be very familiar with unittest. This example uses both to some extent.
  • Lastly, there are many heated debates around testing, what kind, when, where and if. I feel very strongly that testing is necessary in "production" code. I believe there is a delicate balance between unit tests and integration tests.

I am here to act as a guide and hopefully provide some hints for your own journey.

The application

I have put together a very small application, but written in a way that you might encounter on a larger project (not the typical monolithic file you commonly find in tutorials). I will be using this repository throughout this post. We will use yet another messaging app, but this one has much less functionality than other examples you have seen. There is a single GET method (show_messages) that returns a list of messages mapped to an authenticated user. The user object is attached to Flask.g when the login_required decorator fires.

The app leverages an app factory, a Blueprint, and a back-end database running SQLite in memory. The database provider is not relevant here since we are focused on interaction with the ORM, not what is under the covers. The app factory method is located in src/app.py.

The model

The app uses a very simple model. We have users and messages. A message is related to a user by the username foreign key. This is vanilla SQLAlcehmy.

class User(DB.Model):
    """Model representing a user object."""
    username = DB.Column(DB.String, primary_key=True)
    password = DB.Column(DB.String, nullable=False)
    nick = DB.Column(DB.String, nullable=False)

class Message(DB.Model):
    """Message model."""
    message_id = DB.Column(DB.Integer, primary_key=True, autoincrement=True)
    username = DB.Column(DB.String, DB.ForeignKey(
        f"{User.__tablename__}.username"
    ), nullable=False)
    message_text = DB.Column(DB.String, nullable=False)

The message view

The interesting bit of code in is located in the src/message_bp.py file. The show_messages method could be implemented more easily using an extension. Please note the authentication decorator that enforces who can view which posts src/auth.py. The Flask.g object stores the current user in the request context, queried from the database upon login. This is wrapped up in a Response object and sent back to the client.

BP = Blueprint("app", __name__)

@BP.route("/messages")
@AUTH.login_required
def show_messages():
    """Show messages by user."""
    try:
        messages = m.Message.query.filter(
            m.Message.username == g.user.username
        ).all()
    except sqlalchemy.exc.SQLAlchemyError:
        abort(500, "Could not retrieve records from the DB.")
    resp_text = json.dumps([
        {
            "user": message.username,
            "message_text": message.message_text
        }
        for message in messages
    ])
    return Response(resp_text, mimetype="application/json")

Testing with Flask and PyTest - intro

I first learned unittest, and fell in love with PyTest the first time I used it. This code could easily be adapted to unittest.

Flask provides a nice testing mini-framework that we will utilize along the way. Keep in mind that the techniques required to test Flask also apply to any of its extensions. The main "gotcha" that fellow developers run into is when dealing with an applicaiton or request context. The simplest way to deal with this is to use Flask's built-in test_client. Let's take a stab at testing our show_messages method!

First attempt - integration testing

The easist, fastest way to get started with testing your Flask application is have all of the components in place and test the entire system. Many developers start the coding process with a test database that contains replicas of production data. This makes debugging as you code pretty simple and has the advantage of working with "real" objects. You might even have an SQL script that sets up tables in a fresh DB instance before the tests are run. I have written a Python script that uses the ORM to load our test objects into the SQLite instance. It is located at tests/load_db_data.py. Our test runner lives in tests/test_messages.py. I wrote a fixture (similar to unittest's setUp) that loads the DB and creates and application instance. This fixture is passed to each of the test methods. Here's our first test:

def test_get_message_200(app_inst):
    """Send a GET with good auth and expect a message."""
    client = app_inst.test_client()
    auth_str = base64.b64encode(b"foo:foo_pass").decode("utf-8")
    rv = client.get("/app/messages", headers={"Authorization": f"Basic {auth_str}"})
    assert rv.status_code == 200
    assert b'"user": "foo"' in rv.data

The test_client is provided by Flask. I also like using requests for integration tests with larger projects. The auth string is built by base64-encoding a known username and password (setup in the init_db method). These are wrapped up and sent to the application instance and we make assertions based on known database state. That last part is critical because the test cannot run successfully without the database instance existing, and more importantly, containing data relevant to the test.

$ python -m pytest tests

platform darwin -- Python 3.7.0, pytest-3.6.3, py-1.5.4, pluggy-0.6.0
rootdir: pydocker, inifile:
collected 2 items

tests/test_messages.py ..                                                                                                                                                                                                              [100%]

2 passed in 0.30 seconds

Requiring a certain database state (or any other external state) is the Achiles heel of integration testing. We are really interested in testing our business logic, right? We trust that the core developers of Flask and SQLAlchemy have tested their modules; we're more interestd in testing our own business logic. Also, if you are using an on-disk database that lives on a VM that sits in a cloud that sits on spinning hard disks... your tests will be slow! You don't want to shower the Twitter API with automated test calls, either. How can we avoid these problems, is there a better way? Yes, but before we look, remember that integration tests are a critical component of the overall development process! They have their place and it is up to you and your team to determine which method is most appropriate for each situation.

The dreaded unit test

Dreaded? Yeah, I certainly think so - but I hope you see the light before we are done. Writing unit tests with Flask and SQLAlchemy means that you have to decouple application test state and dependencies such as the database connection or calls to 3rd-party API's. This is a non-trivial task that requires a working knowledge of mocking and patches. As a prerequisite, I highly recommend this video from PyCon 2018 by Lisa Roach.

If you refer back to the show_messages method above, you will see a couple issues right off the bat. First and ugliest is that login_required decorator. The next issue is the database query. Heck, let's throw some code at it and see what happens! I'll monkeypatch the ORM object as well as the decorator:

@patch("src.message_bp.m.Message")
@patch("src.message_bp.AUTH", lambda x: x)
def test_show_messages_mock(mock_m, app_inst):
    """Send a GET with good auth and expect a message."""
    mock_m.query.filter.return_value.all.return_value = ["foo", "bar"]
    with app_inst.test_request_context("/app/messages"):
        # Calling the method directly returns a response object
        resp = show_messages()
        assert resp.status_code == 200

python -m pytest tests/ results in:

mock_m = <MagicMock name='Message' id='4350213424'>, app_inst = <Flask 'src.app'>

...
>           assert resp.status_code == 200
E           assert 401 == 200
E            +  where 401 = <Response 19 bytes [401 UNAUTHORIZED]>.status_code

tests/test_messages.py:47: AssertionError

What happened? The problem with the decorator object is that it evaluated at import time and we're trying to patch it at run time. Nuts, that will be kind of hard to work around, right? This is where I would like to make my first major point.

-> When refactoring code during tests or using TDD, think of ways to make your code testable. Constantly ask yourself if you can break things into functional pieces or if you need to strip out dependent code.

How can we do that to show_messages? The first thing that comes to mind is to gut the entire method and move it to its own function. Here's our new version:

def get_messages():
    """Retrieve messages from the DB by user."""
    try:
        messages = m.Message.query.filter(
            m.Message.username == g.user.username
        ).all()
    except sqlalchemy.exc.SQLAlchemyError:
        abort(500, "Could not retrieve records from the DB.")
    resp_text = json.dumps([
        {
            "user": message.username,
            "message_text": message.message_text
        }
        for message in messages
    ])
    return resp_text

@BP.route("/messages")
@AUTH.login_required
def show_messages():
    """Show messages by user."""
    resp_text = get_messages()
    return Response(resp_text, mimetype="application/json")

At this point, we can still leverage our previous integration tests to make sure we didn't break anything. Here is where we make our first testing detour. Now, show_messages lends itself well to an integration test, but it's so simple and concise that we really don't have anything to test as a "unit"... except for that new method we created. It doesn't have any decorator magic that we need to hack around, so let's see if it can be tested by itself. The new test function:

@patch("src.message_bp.m")
def test_show_messages_mock(mock_m, app_inst):
    """Send a GET with good auth and expect a message."""
    # This could be done in a much better way, keep reading!
    mock_m.Message.query.filter.return_value.all.return_value = [
        m.Message(username="test", message_text="test_text"),
        m.Message(username="test", message_text="test_text2"),
        m.Message(username="test2", message_text="test_text"),
        m.Message(username="test2", message_text="test_text2"),
    ]
    resp = get_messages()
    # Now, we're checking for actual data instead of messing with response codes.
    assert json.loads(resp) == [
        {"user": "test", "message_text": "test_text"},
        {"user": "test", "message_text": "test_text2"},
        {"user": "test2", "message_text": "test_text"},
        {"user": "test2", "message_text": "test_text2"}
    ]

If we give that a shot, where does that leave us?

src/message_bp.py:14: in get_messages
    m.Message.username == g.user.username
...
self = <flask.g of 'src.app'>, name = 'user'

    def __getattr__(self, name):
        if name == '__members__':
            return dir(self._get_current_object())
>       return getattr(self._get_current_object(), name)
E       AttributeError: '_AppCtxGlobals' object has no attribute 'user'

Well, snap, the pesky g isn't playing nice now. That's actually an easy fix! Since we are setting up a test context, all we have to do is manually set something on that object. The query itself is being mocked, so it's value has no use to us. Another key point:

-> You cannot call SQLAlchemy queries or modifiy Flask.g outside of a request or application context. This applies even if you are mocking other dependent methods.

...
    # You *must* use an application or request context when dealing with Flask.g !!!
    with app_inst.test_request_context("/app/messages"):
        g.user = MagicMock(username="not a user")
        resp = get_messages()
    ...

I personally believe working through this kind of refactoring is a great exercise. I think the best tests are ones that focus on small chunks of code. Also, never forget what you want to test. If you want to test business logic, then test business logic! For example, in the above messages view, what if you needed to perform a calculation or combine that data with something from another source? I would even factor out the DB call to its own method that simply return a list of (possibly serialized) objects. Methods that return simple objects (even ORM objects) are really quite easy to factor out and test. I'll come back to this after...

A shout-out to Faker!

Do you ever get tired of coming up with mock data? Then check out Faker! I cringed a little when I wrote the mock data above, but I wanted to use it as a pet example. You can easily refactor it to something like this:

mock_m.Message.query.filter.return_value.all.return_value = [
        m.Message(username=FAKER.word(), message_text=FAKER.text())
            for number in range(10000)
    ]

Isn't that so much better?! In the final version of this code, I create 10k objects in the unit test and in the integration tests. Even with an in-memory SQLite database, the unit tests run about 3 times faster due to the overhead of the DB. This will only get worse as your project scales or with more complex providers.

Unit test finale

Finally, I put the query in a standalone function and try to come up with better names.

def query_message_by_user():
    """Retrieve messages from the DB by user."""
    try:
        message_objs = m.Message.query.filter(
            m.Message.username == g.user.username
        ).all()
    except sqlalchemy.exc.SQLAlchemyError:
        abort(500, "Could not retrieve records from the DB.")
    return message_objs

def retrieve_message_text():
    """Retrieve messages from database and return JSON string."""
    messages = query_message_by_user()
    ...

That requires a slight change to our unit test patch.

# Note that we're patching the `query_message...` method and not the DB call.
@patch("src.message_bp.query_message_by_user")
def test_retrieve_message_text(mock_q, app_inst):
    """Send a GET with good auth and expect a message."""
    mock_q.return_value = [
        m.Message(username=FAKER.word(), message_text=FAKER.text())
            for number in range(10000)
    ]
    with app_inst.test_request_context("/app/messages"):
        # You *must* use an application or request context when dealing with Flask.g !!!
        g.user = MagicMock(username="not a user")
        resp = retrieve_message_text()
    assert json.loads(resp)

You may be wondering why make such a small change. The retrieve_message_text is now streamlined down to its core logic. This is a pet example, but this is where that critical business logic should live. It also frees you to deal with DB queries and error handling without polluting this method with exception-handling.

In conclusion

This post assumes a lot of prior knowledge in an attempt to make a subtle point. My goal was to focus on the process of refactoring your code to make it more testable. The techniques are the same regardless of the test framework you are using in Python. Flask and SQLAlchemy present some unique challenges of their own and sometimes require delicate handling. In closing, if you are struggling with a mock or return data, take a step back and see if you can factor the troublesome code into another method. At that point you might be able to simply patch the method instead of the data itself (as in the last example). I have learned the hard way that testing is much more art than science. I hope this helps, and feel free to leave comments on my repo!