from __future__ import annotations import ssl import threading import time from collections.abc import Callable, Generator, Mapping from contextlib import contextmanager from dataclasses import dataclass from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from queue import SimpleQueue from types import MappingProxyType from typing import Final @dataclass(frozen=True, slots=True) class Request: method: str target: str headers: Mapping[str, str] body: bytes @dataclass(frozen=True, slots=True) class Reply: status: int = 200 body: bytes = b"{}" content_type: str = "application/json" chunks: tuple[bytes, ...] | None = None abort_after: int | None = None gate_after_first: threading.Event | None = None pause_between_chunks: float = 0 headers: Mapping[str, str] = MappingProxyType({}) drop_connection: bool = False @dataclass(frozen=True, slots=True) class Wire: url: str received: SimpleQueue[Request] disconnected: SimpleQueue[str] connected: SimpleQueue[str] def drain(self) -> tuple[Request, ...]: return tuple(self.received.get_nowait() for _ in range(self.received.qsize())) def connections(self) -> int: return self.connected.qsize() @contextmanager def wire_server( respond: Callable[[Request], Reply], tls: ssl.SSLContext | None = None, port: int = 0, keep_alive: bool = False, ) -> Generator[Wire, None, None]: """Owned TCP peer; requests traverse the real HTTP client and serialization. With `keep_alive` the peer honours HTTP/1.1 persistent connections so `connections()` counts the client's TCP sessions.""" received: Final[SimpleQueue[Request]] = SimpleQueue() errors: Final[SimpleQueue[Exception]] = SimpleQueue() disconnected: Final[SimpleQueue[str]] = SimpleQueue() connected: Final[SimpleQueue[str]] = SimpleQueue() class Handler(BaseHTTPRequestHandler): protocol_version = "HTTP/1.1" timeout = 5 def setup(self) -> None: super().setup() connected.put(f"{self.client_address[0]}:{self.client_address[1]}") def respond(self) -> None: request: Final = Request( self.command, self.path, {name.lower(): value for name, value in self.headers.items()}, self.rfile.read(int(self.headers.get("content-length", "0"))), ) received.put(request) try: reply = respond(request) except Exception as error: errors.put(error) reply = Reply(status=500) if reply.drop_connection: self.close_connection = True disconnected.put(request.target) return self.send_response(reply.status) self.send_header("content-type", reply.content_type) for name, value in reply.headers.items(): self.send_header(name, value) if reply.chunks is None: self.send_header("content-length", str(len(reply.body))) else: self.send_header("transfer-encoding", "chunked") if not keep_alive: self.send_header("connection", "close") self.end_headers() try: if self.command == "HEAD": self.wfile.flush() elif reply.chunks is None: self.wfile.write(reply.body) else: for index, chunk in enumerate(reply.chunks): if reply.abort_after == index: break self.wfile.write(b"%x\r\n%s\r\n" % (len(chunk), chunk)) self.wfile.flush() if index == 0 and reply.gate_after_first is not None: assert reply.gate_after_first.wait(timeout=5), "Stream barrier was never released" if reply.pause_between_chunks and index + 1 < len(reply.chunks): time.sleep(reply.pause_between_chunks) else: self.wfile.write(b"0\r\n\r\n") self.wfile.flush() except (BrokenPipeError, ConnectionResetError): disconnected.put(request.target) except Exception as error: errors.put(error) self.close_connection = not keep_alive do_POST = respond do_PUT = respond do_GET = respond do_DELETE = respond do_PATCH = respond do_HEAD = respond def log_message(self, format: str, *args: object) -> None: pass class OwnedHTTPServer(ThreadingHTTPServer): daemon_threads = False request_queue_size = 128 def server_bind(self) -> None: super().server_bind() if tls is not None: self.socket = tls.wrap_socket(self.socket, server_side=True) with OwnedHTTPServer(("127.0.0.1", port), Handler) as server: thread: Final = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.05}) thread.start() try: yield Wire( f"{'https' if tls is not None else 'http'}://127.0.0.1:{server.server_port}", received, disconnected, connected, ) finally: server.shutdown() thread.join(timeout=6) assert not thread.is_alive(), "Owned HTTP server survived cleanup" server.server_close() failure: Final = None if errors.empty() else errors.get_nowait() assert failure is None, f"Owned HTTP peer failed: {failure!r}"