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 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):
    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)
        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):
    """
    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),
            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()