An "Easier" way to test Flask authentication?
Recently, I was on the testing train (for millionth time), walking towards the last carriages, trying to add and improve tests for one of my work projects when I stumbled upon this:
# def test_create_application(client: FlaskClient, mocker):
# user = User.create(email='test@gmail.com', password=hash_password('test'))
# with client:
# client.post('/login', data=dict(email=user.email, password='test'))
#
# res = client.post('/create_application', data=dict(name='test', identifier='http://test.com'),
# follow_redirects=True)
# assert res.ok
# assert b'Application added successfully' in res.data
#
# # Test duplicate identifiers
# res = client.post('/create_application', data=dict(name='test', identifier='http://test.com'),
# follow_redirects=True)
# assert b'Identifier must be unique' in res.data
Let me say that I was bothered. Not by the commented out tests (yes, plural; the entire file actually 🙆♂️ ). Not by the fact that there were multiple assertions in a single test function. But by the sudden surge of memories of figuring out how to test flask views which are protected by authentication.
Heh! The stories I can tell!
Anyway, back to the code. One might say it's quite straight forward.
- A user is created
- We use the created user to login
- We make some calls to specific views and do assertions with responses
Now, using the above mentioned steps should work for most (if not all) protected views.
Note: By views, we are referring to non-API routes i.e pages a user would interact within a browser.
But, there's more to this than meets the eye, especially in step 2.
For this web app, users would log in using an email and password. Under the hood, this would involve validating the user's credentials then modifying the current session by saving user information such as a unique user identifier or email. In a different web app, it might involve some form of external authentication. For instance, using an Oauth provider such as Google or Facebook or some kind of Lightweight Directory Access Protocol (LDAP) authentication.
Thinking about all this, one begins to realize that step 2 might not be a one-liner. For instance, in one of our apps, a user logs in using OpenID. We are planning to make this the default way of authentication for all our users.
Is there a "simple" way to approach this? A better way to tell your app "Hey, for the next number of x route calls, assume that the user is y".
This scenario, made me think about testing in Django. Yes, yes, I know that word is a taboo in the flask world 😆. Anyway, in "the name that must not be mentioned" test suite, there's a login method that allows you to log in a user without calling any URLs.
What if? Hear me out. What if we could "borrow" that functionality 🌚 . Take a look at the following sample web app:
from flask import Flask, request, redirect, url_for
from flask_login import LoginManager, UserMixin, login_required, login_user, logout_user
from werkzeug.exceptions import MethodNotAllowed, Unauthorized
class Config:
SECRET_KEY = 'Some very long secret key'
class User(UserMixin):
pass
login_manager = LoginManager()
@login_manager.user_loader
def user_loader(email):
print('[[')
if email != 'test@gmail.com':
return
user = User()
user.id = email
return user
@login_manager.request_loader
def request_loader(request):
email = request.form.get('email')
if email != 'test@gmail.com':
return
user = User()
user.id = email
user.is_authenticated = request.form['password'] == 'test'
return user
@login_manager.unauthorized_handler
def unauthorized_handler():
raise Unauthorized()
def create_app():
app = Flask(__name__)
app.config.from_object(Config)
login_manager.init_app(app)
@app.route('/')
@login_required
def index():
return 'Ok'
@app.route('/login', methods=['GET', 'POST'])
def login():
if request.method != 'POST':
raise MethodNotAllowed()
if request.form['email'] == 'test@gmail.com' and request.form['password'] == 'test':
user = User()
user.id = request.form['email']
login_user(user)
return redirect(url_for('.index'))
raise Unauthorized()
@app.route('/logout')
def logout():
logout_user()
return 'Logged out'
return app
We have a basic flask application with three routes:
- One to log in
- The root page, which is protected
- And one for logging out
For the sake of this article, we are making the app as basic as possible. There are a lot of "best practices " violations 😨, bear with me. 😅
Here are the tests for the app:
import pytest
from flask.testing import FlaskClient
from app import create_app
@pytest.fixture(scope='module')
def flask_app():
app = create_app()
with app.app_context():
yield app
@pytest.fixture(scope='module')
def client(flask_app):
app = flask_app
ctx = flask_app.test_request_context()
ctx.push()
app.test_client_class = FlaskClient
return app.test_client()
def test_index_page__not_logged_in(client):
res = client.get('/')
assert res.status_code == 401
def test_index_page__logged_in(client):
with client:
client.post('/login', data=dict(email='test@gmail.com', password='test'))
res = client.get('/')
assert res.status_code == 200
We have two tests:
- One verifies that an unauthenticated user cannot access the root page.
- The other verifies that an authenticated can access the root page.
We are more interested in the second test.
What if we can make the second test something like:
def test_index_page__logged_in(client):
res = client.get('/')
assert res.status_code == 200
Where we do not have any of the login logic inside the test.
Well, good news, we have decorators in python.
If you are not familiar with decorators in python, feel free to take a look at the following article.
Let's introduce a decorator called force_login. Take a look at the following snippet:
def force_login(email=None):
def inner(f):
@functools.wraps(f)
def wrapper(*args, **kwargs):
if email:
for key, val in kwargs.items():
if isinstance(val, FlaskClient):
with val:
with val.session_transaction() as sess:
sess['_user_id'] = email
return f(*args, **kwargs)
return f(*args, **kwargs)
return wrapper
return inner
We are interested in the following lines:
with val:
with val.session_transaction() as sess:
sess['_user_id'] = email
Since we are using the Flask-login extension, we need to know how it tracks logged in users. Taking a look at the source code of the login_user we can see that it stores a couple of variables in the session:
def login_user(user, remember=False, duration=None, force=False, fresh=True):
...
if not force and not user.is_active:
return False
user_id = getattr(user, current_app.login_manager.id_attribute)()
session['_user_id'] = user_id # This is what we need
session['_fresh'] = fresh
session['_id'] = current_app.login_manager._session_identifier_generator()
Thus the use of sess['_user_id'] = email in our decorator. If we were to use a different approach, such as a custom implementation or a different extension, we would have to adapt our decorator accordingly. Remember, we are not interested in where the user info is coming from. We just want to move our flask app to an authenticated state and do some assertions based on that state.
With the above decorator we can make changes to our second test:
@force_login(email='test@gmail.com')
def test_index_page__logged_in(client):
res = client.get('/')
assert res.status_code == 200
Noice! Definitely, looks cleaner.
We have moved our login logic out of test into the force_login decorator so that we can focus on the test.
We can make use of the decorator package to cleanup our decorator:
@decorator
def force_login(func, email=None, *args, **kwargs):
for arg in args:
if isinstance(arg, FlaskClient):
with arg:
with arg.session_transaction() as sess:
sess['_user_id'] = email
return func(*args, **kwargs)
return func(*args, **kwargs)
Let's see if we can make our decorator a bit more generic. We can pull out the logic for modifying the session. Let's make the following changes:
def login_user(sess, email):
sess['_user_id'] = email
@decorator
def force_login(func, cb=None, *args, **kwargs):
for arg in args:
if isinstance(arg, FlaskClient):
with arg:
with arg.session_transaction() as sess:
cb(sess)
return func(*args, **kwargs)
return func(*args, **kwargs)
Then we can change the decorator call as follows:
@force_login(cb=lambda s: login_user(s, 'test@gmail.com'))
def test_index_page__logged_in(client):
res = client.get('/')
assert res.status_code == 200
We are wrapping the function responsible for modifying the session in a lambda because we want to be able to change the email parameter per test.
The decorator approach isn't limited to authentication. We can use also use it for data insertions into a DB and other use cases I have no idea about 🤯 .
With the above approach, we gain two things:
- Cleaner code - Our tests look cleaner since they do not have login logic.
- Abstraction - Since we have moved all login logic to a decorator, we are able to use the decorator in multiple tests without worrying about the implementation details of authenticating a user.
Top comments (0)