Testing a socketio Web App written in Python

In this article I will show you how you can test an async Socketio application in Python, where the ASGI server we are running is uvicorn. I will be referring to these tests as integration tests, though depending on who you ask they could be called E2E tests, system tests, slow test etc. What I am referring to is simply testing out the entire “flow” of a socketio event i.e. emitting an event from a client, then receiving it on the web service and for my actual projects interacting with an actual database.

We will be using pytest as our testing framework.

main.py

import socketio
import uvicorn
from pydantic import BaseModel


async def startup():
    print("Starting Application")


sio = socketio.AsyncServer(async_mode="asgi")
app = socketio.ASGIApp(sio, on_startup=startup)


class FooEvent(BaseModel):
    name: str


@sio.on("FOO")
async def foo_event(sid, *args, **kwargs):
    data = FooEvent(**args[0])
    await sio.emit("BAR", {"foo": data.name})


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8080)

Let’s take a look at our socketio app. Which is a very simple web app, that listens to one event FOO and responds with a BAR event. It is just this single file.

conftest.py

The conftest.py file is automatically run by pytest and allows our test modules to access fixtures defined in this file. One of the best features of Pytest is fixtures. Fixture are functions that have re-usable bits of code we can run in our tests, such as static data used by tests.

import asyncio
from typing import Any, AsyncIterator, Awaitable, List, Optional

import pytest
import socketio
import uvicorn
from app import main
from socketio import ASGIApp
from socketio.asyncio_client import AsyncClient

PORT = 8000
LISTENING_IF = "127.0.0.1"
BASE_URL = f"http://{LISTENING_IF}:{PORT}"


class UvicornTestServer(uvicorn.Server):
    def __init__(self, app: ASGIApp = main.app, host: str = LISTENING_IF, port: int = PORT):
        self._startup_done = asyncio.Event()
        self._serve_task: Optional[Awaitable[Any]] = None
        super().__init__(config=uvicorn.Config(app, host=host, port=port))

    async def startup(self) -> None:
        """Override uvicorn startup"""
        await super().startup()
        self.config.setup_event_loop()
        self._startup_done.set()

    async def start_up(self) -> None:
        """Start up server asynchronously"""
        self._serve_task = asyncio.create_task(self.serve())
        await self._startup_done.wait()

    async def tear_down(self) -> None:
        """Shut down server asynchronously"""
        self.should_exit = True
        if self._serve_task:
            await self._serve_task


@pytest.fixture(scope="session")
def event_loop():
    loop = asyncio.get_event_loop()
    yield loop
    loop.close()


@pytest.fixture(autouse=True, scope="session")
async def startup_and_shutdown_server():
    server = UvicornTestServer()
    await server.start_up()
    yield
    await server.tear_down()


@pytest.fixture(scope="session")
async def client() -> AsyncIterator[AsyncClient]:
    sio = socketio.AsyncClient()
    await sio.connect(BASE_URL)
    yield sio
    await sio.disconnect()

Quick Aside FastAPI Testing

Now when testing say a FastAPI application, it has a builtin test client we can use. This means we don’t actually have to spin up a Uvicorn server to test our application. We can simply pretend to send requests to the FastAPI web service and it will handle the routing behind the scenes.

We can do something like this, where httpx is a async HTTP client (think like the requests library).

import pytest
from asgi_lifespan import LifespanManager
from httpx import AsyncClient

from app.main import app



@pytest.fixture()
async def client() -> AsyncIterator[AsyncClient]:
    async with LifespanManager(app):
        async with AsyncClient(app=app, base_url="http://localhost") as client:
            yield client

Then we can use it like so in our tests:

from fastapi import status
from httpx import AsyncClient

@pytest.mark.asyncio
async def test_add_game(client: AsyncClient):
    response = await client.post("/game", json=request_data)
    assert response.status_code == status.HTTP_201_CREATED

However socketio at the moment does not provide us with a test client we can use. So we will start and stop a Uvicorn server and send actual Socketio requests from a Socketio client. There is a Socketio client library we can use to do this, available in Python.

class UvicornTestServer(uvicorn.Server):
    def __init__(self, app: ASGIApp = main.app, host: str = LISTENING_IF, port: int = PORT):
        self._startup_done = asyncio.Event()
        self._serve_task: Optional[Awaitable[Any]] = None
        super().__init__(config=uvicorn.Config(app, host=host, port=port))

    async def startup(self) -> None:
        """Override uvicorn startup"""
        await super().startup()
        self.config.setup_event_loop()
        self._startup_done.set()

    async def start_up(self) -> None:
        """Start up server asynchronously"""
        self._serve_task = asyncio.create_task(self.serve())
        await self._startup_done.wait()

    async def tear_down(self) -> None:
        """Shut down server asynchronously"""
        self.should_exit = True
        if self._serve_task:
            await self._serve_task

This is a test class which we can use to start and stop the Uvicorn server. Note that the class inherits from uvicorn.server, we need to overwrite the startup() method as we want to change the startup a bit.

Before explaining the code above let’s take a look at how we may use it:

@pytest.fixture(scope="session")
def event_loop():
    loop = asyncio.get_event_loop()
    yield loop
    loop.close()


@pytest.fixture(autouse=True, scope="session")
async def startup_and_shutdown_server():
    server = UvicornTestServer()
    await server.start_up()
    yield
    await server.tear_down()

@pytest.fixture(scope="session")
async def client() -> AsyncIterator[AsyncClient]:
    sio = socketio.AsyncClient()
    await sio.connect(BASE_URL)
    yield sio
    await sio.disconnect()

What we have done is created two pytest fixtures, the first simply starts an event loop so we can test async code.

Tangent on asyncio

To test async code with pytest we need to install the pyest-asyncio library. By default this will give us an event_loop fixture that runs on scope of function. So it will start and stop after each test function. However if you want to use fixtures that aren’t of scope function i.e. session or module. Then we need to redefine the event_loop function as we have done in the example above.

Okay back to our code above. The main bit we are interested in is the startup_and_shutdown_server function, here we start the server before all of our tests and due to how yield, you can read more about how yield works here, we will stop our server after all of our tests have run.

This happens automatically without calling the function because of the decorator we have provided @pytest.fixture(autouse=True, scope="session"). Again we are using scope session so that this function isn’t called either for every function (which would slow down our tests). We could’ve set it to module but again if we have multiple test files we don’t want to run this function for every file (module).

Deeper diver into UvicornTestServer

Let’s take a look at the first two methods

  def __init__(self, app: ASGIApp = main.app, host: str = LISTENING_IF, port: int = PORT):
      self._startup_done = asyncio.Event()
      self._serve_task: Optional[Awaitable[Any]] = None
      super().__init__(config=uvicorn.Config(app, host=host, port=port))

  async def startup(self) -> None:
      """Override uvicorn startup"""
      await super().startup()
      self.config.setup_event_loop()
      self._startup_done.set()

The __init__ magic dunder method creates an asyncio event asyncio.Event(). These events are often used to:

An asyncio event can be used to notify multiple asyncio tasks that some event has happened. - https://docs.python.org/3/library/asyncio-sync.html#asyncio.Event

Then we create a variable self._serve_task: Optional[Awaitable[Any]] = None, we will see how this used later. Finally we call the parent calls __init__ method (super().__init__()). This calls the __init__ function of the uvicorn.Server class. We do this to set the uvicorn.Config, which includes our app and which host and port to start the server.

Onto the second method startup this also overwrites a method in the parent class. In fact the first we do is call the parent class’s startup method (await super().startup()). Then we start the event loop ourselves self.config.setup_event_loop(), where our web app will run.

Finally we do self._startup_done.set(), we are setting this event as true i.e. is complete. So any coroutines waiting until this set can be carry on their execution.

Yet another tangent on run() method

Now the parent class does have a run method we could use, which would start the event loop for us. This however won’t work, lets pretend we change startup_and_shutdown_server function too look like this (server.run()).

@pytest.fixture(autouse=True, scope="session")
async def startup_and_shutdown_server():
    server = UvicornTestServer()
    await server.run()
    yield
    await server.tear_down()

We would get the following error RuntimeError: asyncio.run() cannot be called from a running event loop. This because if we take a look at the run method in the parent class it contains something like this line return asyncio.run(self.serve(...)).

This is why we need to write our own code to handle starting the Uvicorn server.

start_up and tear_down

Okay let’s move and take a look at the start_up and tear_down methods

async def start_up(self) -> None:
    self._serve_task = asyncio.create_task(self.serve())
    await self._startup_done.wait()

async def tear_down(self) -> None:
    self.should_exit = True
    if self._serve_task:
        await self._serve_task

Remember these are the two methods we will call in our “startup and shutdown” fixture. The start_up method, creates a task and assigns it to our empty variable from the __init__ method self._serve_task = asyncio.create_task(self.serve()). It calls the serve method to start the Uvicorn server.

Then we await self._startup_done.wait(), this is the event we created earlier. It will wait until the set() function has been called in the in the startup method above.

Now onto the tear_down method where we set the should_exit to true. There is a main_loop method called by our serve method in the parent class. This main_loop calls an on_tick function which returns if self.should_exit is true. So the call chain looks like: serve -> main_loop -> on_tick. When on_tick returns should_exist as true, it exits it main loop:

async def main_loop(self) -> None:
    counter = 0
    should_exit = await self.on_tick(counter)
    while not should_exit:
        counter += 1
        counter = counter % 864000
        await asyncio.sleep(0.1)
        should_exit = await self.on_tick(counter)

Client Fixture

Finally lets take a look at our final fixture, here we create a client that can be used to make requests with socketio. We use a similar technique with yields so we return a socketio client. We will see how this used in one of our tests.

@pytest.fixture(scope="session")
async def client() -> AsyncIterator[AsyncClient]:
    sio = socketio.AsyncClient()
    await sio.connect(BASE_URL)
    yield sio
    await sio.disconnect()

test_room.py

import asyncio

import pytest
from socketio.asyncio_client import AsyncClient


@pytest.mark.asyncio
async def test_success(client: AsyncClient):
    future = asyncio.get_running_loop().create_future()

    @client.on("BAR")
    def _(data):
        future.set_result(data)

    await client.emit("FOO", {"name": "haseeb"})
    await asyncio.wait_for(future, timeout=5.0)
    result = future.result()
    assert result == {"foo": "haseeb"}

Since we need to wait for the FOO event to return a BAR event we use a future to await until we get a response then set the return data in the future

@client.on("BAR")
def _(data):
    future.set_result(data)

We await asyncio.wait_for(future, timeout=5.0) for the future to have data set on it.

That’s it, the code itself is fairly simple once everything is setup in conftest to actually do the test.

Appendix

• Example source code
• Github Issue: UvicornTestServer
• Async Create Task SO
• Real application using this testing pattern