miguelgrinberg.python-socketio/tests/asyncio_web_server.py

import threading
import time

import requests
import uvicorn

import socketio


class SocketIOWebServer:
    """A simple web server used for running Socket.IO servers in tests.

    :param sio: a Socket.IO server instance.

    Note 1: This class is not production-ready and is intended for testing.
    Note 2: This class only supports the "asgi" async_mode.
    """

    def __init__(self, sio, on_shutdown=None):
        if sio.async_mode != "asgi":
            raise ValueError('The async_mode must be "asgi"')

        async def http_app(scope, receive, send):
            await send(
                {
                    "type": "http.response.start",
                    "status": 200,
                    "headers": [("Content-Type", "text/plain")],
                }
            )
            await send({"type": "http.response.body", "body": b"OK"})

        self.sio = sio
        self.app = socketio.ASGIApp(sio, http_app, on_shutdown=on_shutdown)
        self.httpd = None
        self.thread = None

    def start(self, port=8900):
        """Start the web server.

        :param port: the port to listen on. Defaults to 8900.

        The server is started in a background thread.
        """
        self.httpd = uvicorn.Server(config=uvicorn.Config(self.app, port=port))
        self.thread = threading.Thread(target=self.httpd.run)
        self.thread.start()

        # wait for the server to start
        while True:
            try:
                r = requests.get(f"http://localhost:{port}/")
                r.raise_for_status()
                if r.text == "OK":
                    break
            except:
                time.sleep(0.1)

    def stop(self):
        """Stop the web server."""
        self.httpd.should_exit = True
        self.thread.join()
        self.thread = None