Unit Testing Applications that use Flask-Login and Flask-SocketIO

One of the useful features of my Flask-SocketIO extension is the test client, which allows you to write Socket.IO unit tests. A long time limitation of the test client was that it did not see cookies set by Flask, such as the user session. This complicated writing Socket.IO tests for applications that require authentication, because most authentication mechanisms write something to the user session or a custom cookie. The use case that caused pain to a lot of developers was applications that use Flask-Login combined with Flask-SocketIO. To unit test such an application you had to resort to weird tricks such as mocking the current_user variable.

I recently came up with a solution to this problem, so I'm glad to report that this limitation is now a thing of the past. In this short article I want to show you how to set up your project to take advantage of the new cookie support in the Socket.IO test client.

An Example Socket.IO Server with Authentication

Let me show you a very simple Socket.IO server that authenticates users:

from flask import Flask, request, abort
from flask_login import LoginManager, login_user, current_user, UserMixin
from flask_socketio import SocketIO, emit

allowed_users = {
    'foo': 'bar',
    'python': 'is-great!',
}

app = Flask(__name__)
app.config['SECRET_KEY'] = 'top secret!'

login = LoginManager(app)
socketio = SocketIO(app)

@login.user_loader
def user_loader(id):
    return User(id)

class User(UserMixin):
    def __init__(self, username):
        self.id = username

@app.route('/login', methods=['POST'])
def login():
    username = request.form['username']
    password = request.form['password']
    if username not in allowed_users or allowed_users[username] != password:
        abort(401)
    login_user(User(username))
    return ''

@socketio.on('connect')
def on_connect():
    if current_user.is_anonymous:
        return False
    emit('welcome', {'username': current_user.id})

if __name__ == '__main__':
    socketio.run(app)

This is an extremely stripped down server with just the necessary to demonstrate how to write a unit test. The /login route is where the client sends the POST request that logs the user in. The database of users in this example is stored in the allowed_users dictionary to keep things simple.

The interesting part is the connect event handler, where the current_user variable is invoked to check if the client logged in before attempting to connect via Socket.IO. If the client is logged in (i.e. current_user is set to a non-anonymous user), then the handler emits a welcome event back to the client, and includes the username of the logged in user as data. If the user is not logged in, then the Socket.IO connection is rejected by returning False, and nothing is emitted back.

Unit Testing the Application

Without further ado, here is a unit test that exercises the application from the previous section:

from my_app import app, socketio

def socketio_test():
    # log the user in through Flask test client
    flask_test_client = app.test_client()

    # connect to Socket.IO without being logged in
    socketio_test_client = socketio.test_client(
        app, flask_test_client=flask_test_client)

    # make sure the server rejected the connection
    assert not socketio_test_client.is_connected()

    # log in via HTTP
    r = flask_test_client.post('/login', data={
        'username': 'python', 'password': 'is-great!'})
    assert r.status_code == 200

    # connect to Socket.IO again, but now as a logged in user
    socketio_test_client = socketio.test_client(
        app, flask_test_client=flask_test_client)

    # make sure the server accepted the connection
    r = socketio_test_client.get_received()
    assert len(r) == 1
    assert r[0]['name'] == 'welcome'
    assert len(r[0]['args']) == 1
    assert r[0]['args'][0] == {'username': 'python'}


if __name__ == '__main__':
    socketio_test()

As you can see, this test uses the test clients from Flask and Flask-SocketIO. Both clients are needed because to properly test this application we need to make HTTP and Socket.IO calls. The new feature that enables Socket.IO to see Flask cookies and user session is the flask_test_client argument passed when creating the test client. When this argument is passed, the Socket.IO test client is going to import any cookies that exist in the Flask application.

The first time the socketio_test_client is created the user is not logged in, so the connection fails. The test ensures that the server did not accept the connection.

Next the Flask test client is used to log in to the application by sending a POST request to /login, and passing one of the known username/password combinations. The login route will invoke the login_user() function from Flask-Login, which in turn will record the logged in user in the Flask user session.

Then the test creates a new Socket.IO connection through the test client, and this time it makes sure that the server responded with the welcome event and included the username as an argument.

And that's it, really. All you need to remember is to link the two test clients by passing the Flask test client as an argument to the Socket.IO test client!