Flask, flask_login, pytest: How do I set flask_login's current_user?

Viewed 615

I am trying to use pytest to unit test my Flask app. I have the following test case for an endpoint that requires information from flask_login's current_user:

def test_approval_logic():
    with app.test_client() as test_client:
        app_url_put = '/requests/process/2222'

        with app.app_context():
            user = User.query.filter_by(uid='xxxxxxx').first()
            with app.test_request_context():
                login_user(user)
                user.authenticated = True
                db.session.add(user)

                data = dict(
                    state='EXAMPLE_STATE_NAME',
                    action='approve'
                )
                resp = test_client.put(app_url_put, data=data)
                assert resp.status_code == 200

Inside the test_request_context, I am able to set current_user correctly. However, this test fails because in the requests view where the PUT is handled, there is no logged in user and 500 error results. The error message is, AttributeError: 'AnonymousUserMixin' object has no attribute 'email'. Can someone explain why current_user goes away and how I can set it correctly?

3 Answers

Using test client to dispatch a request

The current session is not bound to test_client, so the request uses a new session.

Set the session cookie on the client, so Flask can load the same session for the request:

from flask import session

def set_session_cookie(client):
    val = app.session_interface.get_signing_serializer(app).dumps(dict(session))
    client.set_cookie('localhost', app.session_cookie_name, val)

Usage:

# with app.test_client() as test_client:                            # Change these
#     with app.app_context():                                       #
#         with app.test_request_context():                          #
with app.test_request_context(), app.test_client() as test_client:  # to this
    login_user(user)
    user.authenticated = True
    db.session.add(user)

    data = dict(
        state='EXAMPLE_STATE_NAME',
        action='approve'
    )
    set_session_cookie(test_client)  # Add this
    resp = test_client.put(app_url_put, data=data)

About the compatibility of with app.test_request_context()

i. with app.test_client()

with app.test_client() preserves the context of requests (Flask doc: Keeping the Context Around), so you would get this error when exiting an inner with app.test_request_context():

AssertionError: Popped wrong request context. (<RequestContext 'http://localhost/requests/process/2222' [PUT] of app> instead of <RequestContext 'http://localhost/' [GET] of app>)

Instead, enter app.test_request_context() before app.test_client() as shown above.

ii. with app.app_context()

with app.test_request_context() already pushes an app context, so with app.app_context() is unnecessary.

Using test request context without dispatching a request

From https://flask.palletsprojects.com/en/2.0.x/api/#flask.Flask.test_request_context:

This is mostly useful during testing, where you may want to run a function that uses request data without dispatching a full request.

Usage:

data = dict(
    state='EXAMPLE_STATE_NAME',
    action='approve'
)
with app.test_request_context(data=data):  # Pass data here
    login_user(user)
    user.authenticated = True
    db.session.add(user)

    requests_process(2222)  # Call function for '/requests/process/2222' directly

My guess is that no session cookie is passed in your PUT request.

Here is an example of how I log a user during my tests (I personally user unittest instead of pytest, so I reduced the code to the strict minimum, but let me know if you want a complete example with unittest)

from whereyourappisdefined import application
from models import User
from flask_login import login_user

# Specific route to log an user during tests
@application.route('/auto_login/<user_id>')
def auto_login(user_id):
    user = User.query.filter(User.id == user_id).first()
    login_user(user)
    return "ok"

def yourtest():
    application.config['TESTING'] = True # see my side note
    test_client = application.test_client()
    response = test_client.get(f"/auto_login/1")

    app_url_put = '/requests/process/2222'
    data = dict(
        state='EXAMPLE_STATE_NAME',
        action='approve'
    )
    r = test_client.put(app_url_put, data=data)
    

In the documentation we can read: https://werkzeug.palletsprojects.com/en/2.0.x/test/#werkzeug.test.Client

The use_cookies parameter indicates whether cookies should be stored and sent for subsequent requests. This is True by default but passing False will disable this behavior.

So during the first request GET /auto_login/1 the application will receive a session cookie and keep it for further HTTP requests.

Side note:

During setup, the TESTING config flag is activated. What this does is disable error catching during request handling so that you get better error reports when performing test requests against the application.

Here's how I do it on my sites:

user = User.query.filter_by(user_id='xxxxxxx').one_or_none()
if user:
    user.authenticated = True
    db.session.add(user)
    db.session.commit()
    login_user(user)
else:
   # here I redirect to an unauthorized page, as the user wasn't found

I don't know the order is the issue or just the absence of db.session.commit(), but I think you need to have done both in order for your put request to work.

Note, also, that I am using a one_or_none() because there shouldn't be a possibility of multiple users with the same user_id, just a True or False depending on whether a user was found or not.

Related