Testing HTTP Basic Auth in Flask

I created a quick & dirty admin page for a flask-based website I'm working on, and I had to add HTTP authentication support for it. Doing this is easy.

In the app, I add:

import auth

# ...

@app.route('/admin/', methods=['GET'])
@auth.requires_auth
def admin_page():
    # ... code to fetch info from db ...
    return render_template('admin.html',
        # ... args ...
        )

Note the @auth.requires_auth decorator. This uses an auth module (code stolen from this snippet):

from functools import wraps
from flask import request, Response

def check_auth(username, password):
    """This function is called to check if a username /
    password combination is valid.
    """
    return username == 'me' and password == 'my_password'

def authenticate():
    """Sends a 401 response that enables basic auth"""
    return Response(
    'Could not verify your access level for that URL.\n'
    'You have to login with proper credentials', 401,
    {'WWW-Authenticate': 'Basic realm="Login Required"'})

def requires_auth(f):
    @wraps(f)
    def decorated(*args, **kwargs):
        auth = request.authorization
        if not auth or not check_auth(auth.username, auth.password):
            return authenticate()
        return f(*args, **kwargs)
    return decorated

Easy enough to do a one-off test, but it really ought to have an automated test. Unfortunately, I don't see support for setting headers in flask's built-in test client or in the Flask-Testing extension. That's ok, we can use the test client in Werkzeug.

First, the imports we need:

from flask.ext.testing import TestCase
from werkzeug.test import Client
from werkzeug.datastructures import Headers
import myapp

Note that the TestCase class we're using requires Flask-Testing (pip install Flask-Testing; I'm using version 0.4). This test class wants us to set up the app:

class AuthTest(TestCase):

    def create_app(self):
        views.DATABASE = 'testing.sqlite3'
        if os.path.exists(views.DATABASE):
            os.unlink(views.DATABASE)
        query_db(self.CREATE_TABLE)
        return myapp.app

Then writing a test case to verify that the admin page is locked is straightforward:

    def test_admin_page_is_locked(self):
        rv = self.client.get('/admin/')
        self.assert_401(rv)
        self.assertTrue('WWW-Authenticate' in rv.headers)
        self.assertTrue('Basic' in rv.headers['WWW-Authenticate'])

But actually sending an Authorization header for the client is a bit more challenging. Note that self.client is a subclass of Werkzeug's test.Client, and the open method there supports passing in extra Werkzeug Header objects. This is ugly, but it works. The first two cases below test login attempts with bad password and username, respectively, and the final test case logs in with the correct credentials.

    def test_admin_page_rejects_bad_password(self):
        h = Headers()
        h.add('Authorization', 'Basic ' + base64.b64encode('me:foo'))
        rv = Client.open(self.client, path='/admin/',
                         headers=h)
        self.assert_401(rv)

    def test_admin_page_rejects_bad_username(self):
        h = Headers()
        h.add('Authorization',
              'Basic ' + base64.b64encode('foo:my_password))
        rv = Client.open(self.client, path='/admin/',
                         headers=h)
        self.assert_401(rv)

    def test_admin_page_allows_valid_login(self):
        h = Headers()
        h.add('Authorization',
              'Basic ' + base64.b64encode(me:my_password))
        rv = Client.open(self.client, path='/admin/',
                         headers=h)
        self.assert_200(rv)

If I have some time when I'm done with this project, I'll probably send a patch to Flask-Testing to add support for headers= to that test client.

Posted on 2013-01-30 by brian in tools .
Comments on this post are closed. If you have something to share, please send me email.