Testing User Sessions with Flask Elegantly

There is excellent information available on testing Flask e.g. http://flask.pocoo.org/docs/0.12/testing/

One thing I ran into was the need to set up a user session for most of my test cases.
The documentation suggests something like:

with app.test_client() as c:
    with c.session_transaction() as sess:
        sess['a_key'] = 'a value'

but repeating that in every test case (especially with a bit more code to set up the session) is not really DRY.

Now the same documentation suggests a very elegant trick using context managers for setting up users on the global object:

from contextlib import contextmanager
from flask import appcontext_pushed, g

@contextmanager
def user_set(app, user):
    def handler(sender, **kwargs):
        g.user = user
    with appcontext_pushed.connected_to(handler, app):
        yield

Inspired by this, I figured that they can be combined into:

@contextmanager
def user_set(app, user):
   with app.test_client() as c:
       with c.session_transaction() as sess:
           start_user_session(sess, user)
       yield c


and then use it in a test like:

user = ...
with user_set(app, user) as client:
    client.get('/logged_in_resource')

Leave a Reply