lix/tests/functional2/testlib/fixtures/http_server.py

"""
HTTP server fixture for tests which binds to an auto-assigned port on localhost.
"""

import asyncio
import contextlib
import dataclasses
import logging
import ssl
import time
import socket
import threading
from queue import Queue
import aiohttp.web as web
import aiohttp.http_writer as http_writer

logger = logging.getLogger(__name__)


@dataclasses.dataclass
class HttpServer:
    app: web.Application
    port: int


class EventTS(asyncio.Event):
    """
    A thread safe version of the asyncio Event

    NOTE: clear() is not thread safe

    Taken from https://stackoverflow.com/a/33006667
    """

    def __init__(self, *args, loop: asyncio.AbstractEventLoop | None = None, **kwargs):
        """
        Creates a thread-safe event for the given loop (or the loop of the current thread).
        """
        super().__init__(*args, **kwargs)
        self.target_loop = loop or asyncio.get_running_loop()

    def set(self):
        self.target_loop.call_soon_threadsafe(super().set)


def _make_localhost_socket(port: int = 0) -> tuple[socket.socket, int]:
    """Creates a localhost-bound socket with an auto-assigned port (unless specified)."""
    sock = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
    # Shouldn't matter because we dynamically allocate ports, but this is generally preferred.
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    sock.bind(("::1", port))
    _, port = sock.getsockname()[:2]

    return (sock, port)


def _server_thread(
    app: web.Application,
    sock: socket.socket,
    shutdown_ev_q: Queue,
    ssl_context: ssl.SSLContext | None = None,
):
    async def async_main():
        nonlocal app, sock
        # Due to Reasons(tm) of event loop lifecycles and stuff of the sort,
        # it's far easier to just send the event object to the other thread
        # from inside the loop where it already knows which loop it is.
        shutdown_ev = EventTS()
        shutdown_ev_q.put(shutdown_ev)

        runner = web.AppRunner(app, handle_signals=False)
        await runner.setup()
        site = web.SockSite(runner, sock, ssl_context=ssl_context)
        await site.start()
        await shutdown_ev.wait()
        await runner.cleanup()

    asyncio.run(async_main())


@contextlib.contextmanager
def http_server(app: web.Application, port: int = 0, ssl_context: ssl.SSLContext | None = None):
    """
    Creates an http server on an automatically chosen port (if not given) on
    the host running the given web.Application, gives you the port for it.

    The server is run on a separate thread.
    """
    # n.b. pytest doesn't directly support asyncio. There's a bunch of
    # complexity that we could go through to do this or we could just throw the
    # async on a thread which was what we would do to the web server anyway if
    # it was blocking.
    shutdown_ev_q = Queue()
    thr = None
    sock = None
    shutdown_ev = None
    try:
        sock, port = _make_localhost_socket(port=port)
        thr = threading.Thread(
            target=_server_thread,
            args=(app, sock, shutdown_ev_q, ssl_context),
            name=f"functional2 httpd [::1]:{port}",
        )
        thr.start()
        shutdown_ev = shutdown_ev_q.get()
        yield HttpServer(app=app, port=port)
    finally:
        if shutdown_ev:
            shutdown_ev.set()
        if thr:
            thr.join()
        if sock:
            sock.close()


def dev_main():
    """A little test server for poking at this manually"""
    logging.basicConfig(level=logging.DEBUG)

    async def root(_req: web.Request) -> web.Response:
        # sadly required to make this function async
        await asyncio.sleep(0.01)
        return web.Response(body="hello world")

    async def eof(req: web.Request) -> web.StreamResponse:
        # sadly required to make this function async
        await asyncio.sleep(0.01)
        writer: http_writer.StreamWriter = req.writer  # type: ignore
        assert writer.transport
        writer.transport.close()
        return web.StreamResponse()

    async def eof_after_headers(req: web.Request) -> web.StreamResponse:
        resp = web.StreamResponse()
        resp.headers["content-length"] = "1337"
        await resp.prepare(req)
        await resp.write_eof()
        writer: http_writer.StreamWriter = req.writer  # type: ignore
        assert writer.transport
        writer.transport.close()
        return resp

    app = web.Application()
    app.add_routes(
        [web.get("/", root), web.get("/eof", eof), web.get("/eof_after_headers", eof_after_headers)]
    )

    with http_server(app, 1337) as httpd:
        logger.info("Listening on http://[::1]:%d", httpd.port)
        while True:
            time.sleep(3600)


if __name__ == "__main__":
    dev_main()