2019-02-16T19:29:15Z

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!

9 comments

  • #1 Ramazan Polat said 2019-02-17T15:59:52Z

    Another great post, thank you Miguel. Is there any other way to authenticate and maintain the authentication status with only socketio? This post illustrates a good example to recognize a client using cookies but client always needs to post a form and send username and password before have a websocket connection. Is there any way to do it without a REST requrest, only involving socketio connection?

  • #2 Miguel Grinberg said 2019-02-17T23:49:57Z

    @Ramazan: Yes. The reason the example in this article uses a POST request is that many applications use the traditional login form to log users in with Flask-Login. But if you don't have that, then you can authenticate by passing a token or API key.

  • #3 Sarah said 2019-02-22T10:20:22Z

    Hi Miguel, Thank you for the great tutorial! I want to ask if the socketio server would create a thread for each new connection? And,, How to count the connections in a socketio channel? Thank you very much for replying, thank you! Have a nice day!

  • #4 Miguel Grinberg said 2019-02-22T16:41:51Z

    @Sarah: connections are tied to tasks more than threads. Depending on the async mode that you use you will be using threads, or more likely green threads. But as a general statement yes, each connection runs on a separate context. To count the connections you can maintain a counter in the connect and disconnect handlers.

  • #5 Sarah said 2019-03-06T06:22:13Z

    Hi Miguel, Thanks so much for your help!! Could I see ”socketio.run“ as a kind of middleware if WSGI?

  • #6 Miguel Grinberg said 2019-03-06T09:05:23Z

    @Sarah: Not really. The socketio.run() method starts the web server. It isn't really required that you use this function, you could just start your web server directly if you prefer. The nice thing about this method is that it includes the logic to start the eventlet, gevent and Flask web servers and it starts the correct one automatically for you.

  • #7 Sarah said 2019-04-18T04:31:37Z

    Hi Miguel, Thank you for the kind help. I want to ask if the emit are keep in the app.contex ? I want to notify the client the countdown number in a threading.Timer on server, but even if I emit with the app.contex in server, the client doesn’t receive anything. Is the context copied failed? How to use the flask_socketio. Emit in the func_wrapper? I print the emit in the timer thread: <function emit at 0x10c7051e0>, it’s as same as I print in the main thread And.. is there a better method to implement countdown asynchronously? Thank you very much!! Thanks! ``` def notify_browser(app,data, sendroom): with app.app_context(): print("notify_browser emit",emit) emit('countdown', data, namespace = '/test', room= sendroom) def set_interval(the_app, notify_browser, sec, times, sendroom): def func_wrapper(the_app, cntnum): set_interval(the_app,notify_browser, sec, cntnum, sendroom) notify_browser(the_app,cntnum,sendroom) if times==0: return else: times-=1 t = threading.Timer(sec, func_wrapper,[the_app,times]) t.start() return t ```

  • #8 Miguel Grinberg said 2019-04-18T14:47:11Z

    @Sarah: are you using eventlet or gevent? If you are, then the threading package needs to be monkey patched to work well with these frameworks.

  • #9 Sarah said 2019-04-28T07:52:43Z

    Hi Miguel, Thank you for the kind help. I found I don't need to monkey patched if I import threading from eventlet.green! Thanks a lot!!

Leave a Comment