Integration testing websockets with Python

TL;DR - Testing is important but can be difficult. I've written the Python library pywsitest to help with integration testing websockets and included some examples.

At Grid Smarter Cities we believe that testing is a key component of good software.

Our attitude towards testing is reflected primarily in our development process where we practice test driven development, perform code analysis using Prospector, and require 100% test coverage (which we check using Coverage.py). This process is re-iterated in our automated build pipelines, which run linting, the full suite of unit tests against our code, and integration tests against a staging environment.

A lot of our ideas about testing are roughly based on a testing pyramid (see below), where because unit tests are fast to write and fast to run, they make up the majority of tests we use. Integration tests take longer to run, generally incur some cost running against a live system, and are more complicated to write; so we write less of them, but they're still an important part of ensuring we have a good product.

UI tests and manual testing are both more time consuming and more expensive processes, so they make up a smaller portion of our tests. Because they're not something I have to worry too much about as a back-end developer I'm going to skip over them a bit today, though I appreciate the importance of a full testing suite and process.

However many unit tests we write, and how much we can trust our code in isolation, we can't have full confidence in our products unless we also test interactions between modules on a live system. Automated integration testing of REST APIs is a reasonably straight forward process using tools like Dredd alongside OpenAPI templates, but Grid ran into difficulties testing some of our more complicated interactions with a websocket based API.

pywsitest

One of our products that has been developed recently at Grid required different users interacting with and updating a request object in real time. We chose to use a websocket-based API for all of the asynchronous parts of our application.

When each user type connects to the websocket host, they receive a set of messages about their current request status, other active users, and some geolocation-based data. When one user modifies a request object, certain other users receive the status update, and users are restricted as to which actions they can perfom on a request.

Testing asynchronous interations between the users became a challenge, and is what lead us to creating our Python library pywsitest (python websocket integration testing framework).

The pywsitest library allows a user to connect to a websocket host, assert that a series of messages have been received, and that any messages that need to be sent were sent successfully. Messages can either be sent on connection, or be triggered by receiving any specified message. This ability to listen for responses, and send messages out when they're received allows for writing tests that target the interaction between different users.

Testing scenario

I've written a websocket based chat server that allows a user to connect using a username, send messages to all connected users through the server, and disconnect from the server when they're finished.

The user's name will be included in the url as a query parameter. Since I'm running this server locally for testing, the url I'll use to connect will be wss://localhost?name=Rob.

When a user connects to the server, the server will broadcast a message to all connected users to tell them the user has connected:

{ "message": "Rob has connected" }

When a user sends a message to the server, it will be broadcast to all connected users with the format:

{ "message": "Rob: Hello, world!" }

When a user disconnects, the server will broadcast to all connected users:

{ "message": "Rob has disconnected" }

For my first test I want to ensure a user can connect to the websocket host, and assert that they receive a broadcasted message saying that they've connected.

To do this I set up a WSTest object pointing at the correct url, I then added my name as a query parameter. Because we are expecting a message from the host letting us know we've successfully connected, I added an expected response using the with_response method on the WSTest instance. Passed into the with_response method is an instance of WSResponse with a single attribute matching with the json message I'm expecting to receive from the host.

import asyncio
from pywsitest import WSTest, WSResponse

ws_test = (
    WSTest("wss://localhost")
    .with_parameter("name", "Rob")
    .with_response(
        WSResponse()
        .with_attribute("message", "Rob has connected")
    )
)

asyncio.get_event_loop().run_until_complete(ws_test.run())

assert ws_test.is_complete()

As the run method is asynchronous I'm running it synchronously using the asyncio library in this example with the line:

asyncio.get_event_loop().run_until_complete(ws_test.run())

The next functionality I want to test is the user's abililty to send a message after connecting to the websocket host. This test is set-up in a very similar way to the previous test, but with a message triggered when the connected message is received. The test runner will also expect a response from the host displaying the message that the test runner sent.

import asyncio
from pywsitest import WSTest, WSResponse, WSMessage

ws_test = (
    WSTest("wss://localhost")
    .with_parameter("name", "Rob")
    .with_response(
        WSResponse()
        .with_attribute("message", "Rob has connected")
        .with_trigger(
            WSMessage()
            .with_attribute("message", "Hello, world!")
        )
    )
    .with_response(
        WSResponse()
        .with_attribute("message", "Rob: Hello, world!")
    )
)

asyncio.get_event_loop().run_until_complete(ws_test.run())

assert ws_test.is_complete()

The last functionality I want to demonstrate will be testing two clients running simultaneously. To do this I'll create methods to run each of the tests, and use asyncio.gather() to run both of the methods simultaneously. I'll also add a short delay in one of the methods to help ensure that both users have connected before one of them broadcasts a message.

In this test, user_1 will connect to the websocket host, and wait until it receives a broadcasted message that user_2 has connected before sending a message. user_2 will then reply to that message, at which point user_1's test will finish and user_2 will listen for the broadcasted message notifying everyone that user_1 has disconnected.

import asyncio
from pywsitest import WSTest, WSResponse, WSMessage

async def test_user_1():
    ws_test_1 = (
        WSTest("wss://localhost")
        .with_parameter("name", "user_1")
        .with_response(
            WSResponse()
            .with_attribute("message", "user_1 has connected")
        )
        .with_response(
            WSResponse()
            .with_attribute("message", "user_2 has connected")
            .with_trigger(
                WSMessage()
                .with_attribute("message", "Hello from user_1!")
            )
        )
        .with_response(
            WSResponse()
            .with_attribute("message", "user_2: Hello from user_2!")
        )
    )

    await ws_test_1.run()
    assert ws_test_1.is_complete()

async def test_user_2():
    await asyncio.sleep(2)

    ws_test_2 = (
        WSTest("wss://localhost")
        .with_parameter("name", "user_2")
        .with_response(
            WSResponse()
            .with_attribute("message", "user_2 has connected")
        )
        .with_response(
            WSResponse()
            .with_attribute("message", "user_1: Hello from user_1!")
            .with_trigger(
                WSMessage()
                .with_attribute("message", "Hello from user_2!")
            )
        )
        .with_response(
            WSResponse()
            .with_attribute("message", "user_1 has disconnected")
        )
    )

    await ws_test_2.run()
    assert ws_test_2.is_complete()

async def run_tests():
    await asyncio.gather(test_user_1(), test_user_2())

asyncio.get_event_loop().run_until_complete(run_tests())

Closing notes

pywsitest was written to help enable Grid Smarter Cities to better test our products, and simplifies the process of writing integration tests targeting a websocket-based service.

It's an open source project, so if you want to get involved in future development it'd be greatly appreciated!