Flask: How to Test WebSocket with Flask-SocketIO

I had a project built with Flask (Python) and wanted to add some tests.

In the project, there’s a WebSocket using Flask-SocketIO — but how do we test that?

I did some digging and here’s what I found.

Talk is cheap. Show me the code

Start by creating a simple Flask + SocketIO app main.py:

from flask import Flask, render_template
from flask_socketio import SocketIO, send

socketio = SocketIO()

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

@socketio.on("message")
def handle_message(data):
    print("server socketio received message: " + data)
    send(f"server got: {data}")

@socketio.on_error()
def error_handler(e):
    print("socketio error:", e)

@app.route("/")
def hello_world():
    return "Hello, World!"

def create_app():
    socketio.init_app(app)
    return app

if __name__ == "__main__":
    app = create_app()
    socketio.run(app)

Then create a pytest test file test_main.py:

import pytest
from main import create_app, socketio
from flask.testing import FlaskClient
from flask_socketio.test_client import SocketIOTestClient

@pytest.fixture()
def app():
    app = create_app()
    app.config.update({"TESTING": True})
    yield app

@pytest.fixture()
def flask_client(app):
    return app.test_client()

@pytest.fixture()
def runner(app):
    return app.test_cli_runner()

Next, create a SocketIO test client with socketio.test_client(app)

The socketio.test_client(app) must use the app created from the pytest fixture above:

@pytest.fixture()
def socketio_client(app):
    socketio_client = socketio.test_client(app)
    return socketio_client

⚠️ Gotcha!

The socketio in socketio.test_client(app) is the same socketio = SocketIO() from main.py.

Write a test case for SocketIO:

def test_socketio_message(socketio_client: SocketIOTestClient):
    socketio_client.send("hello from pytest")
    recv = socketio_client.get_received()
    print("test_socketio_message recv", recv)

    assert len(recv) == 1
    assert recv[0]["name"] == "message"
    assert recv[0]["args"] == "server got: hello from pytest"
    assert recv[0]["namespace"] == "/"

Then run the test with:

pytest

Done! We can finally test WebSocket! 🎉

Source code

https://github.com/wuttinanhi/flask-socketio-test