Unit Testing Applications that use Flask-Login and Flask-SocketIO
source link: https://www.tuicool.com/articles/hit/JZFzemq
Go to the source link to view the article. You can view the picture content, updated content and better typesetting reading experience. If the link is broken, please click the button below to view the snapshot at that time.
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!
Recommend
About Joyk
Aggregate valuable and interesting links.
Joyk means Joy of geeK