Django Websocket Testing using pytest

async_server.py

"""
This will be in cookiecutter-django so it's easier to copy and understand.
To set up websockets, you should use: https://dev.to/jaydenwindle/adding-websockets-to-your-django-app-with-no-extra-dependencies-2f6h
However, like in cookiecutter-django, you should put the websocket.py inside an app.
It's probably much better to use the stuff in cookiecutter-django (with use_async=y)
since I've already tested this functionality heavily, and those configurations are the best.

"""
import asyncio
import functools
import threading
import time
from contextlib import contextmanager

from uvicorn.config import Config
from uvicorn.main import ServerState
from uvicorn.protocols.http.h11_impl import H11Protocol
from uvicorn.protocols.websockets.websockets_impl import WebSocketProtocol


def run_loop(loop):
    loop.run_forever()
    loop.close()


@contextmanager
def run_server(path="/"):
    asyncio.set_event_loop(None)
    loop = asyncio.new_event_loop()
    config = Config(
        app="config.asgi:application", ws=WebSocketProtocol, http=H11Protocol
    )
    server_state = ServerState()
    protocol = functools.partial(H11Protocol, config=config, server_state=server_state)
    create_server_task = loop.create_server(protocol, host="127.0.0.1")
    server = loop.run_until_complete(create_server_task)
    port = server.sockets[0].getsockname()[1]
    url = f"ws://127.0.0.1:{port}{path}"
    try:
        # Run the event loop in a new thread.
        thread = threading.Thread(target=run_loop, args=[loop])
        thread.start()
        # Return the contextmanager state.
        yield url
    finally:
        # Close the loop from our main thread.
        while server_state.tasks:
            time.sleep(0.01)
        loop.call_soon_threadsafe(loop.stop)
        thread.join()

test_socket.py

"""
Refer to Uvicorn's tests to know how to write your own.
https://github.com/encode/uvicorn/blob/master/tests/protocols/test_websocket.py
"""
from asyncio import new_event_loop

from websockets import connect

from .async_server import run_server


def test_accept_connection():
    """
    If you want to communicate over HTTP, add live_server fixture
    """

    async def open_connection(url):
        async with connect(url) as websocket:
            return websocket.open

    with run_server() as url:
        loop = new_event_loop()
        is_open = loop.run_until_complete(open_connection(url))
        assert is_open
        loop.close()


def test_ping():
    async def ping(url):
        async with connect(url) as websocket:
            await websocket.send("ping")
            return await websocket.recv()

    with run_server() as url:
        loop = new_event_loop()
        received_message = loop.run_until_complete(ping(url))
        assert received_message == "pong"
        loop.close()