"""Shared fixtures + a fake vLLM backend transport (no GPU / docker / network). The router's outbound httpx clients are built with an injectable transport (`ServiceManager(cfg, transport=...)`), so a single fake serves all three backends and records every call the router makes. """ from __future__ import annotations import asyncio import json import sys import urllib.parse from types import SimpleNamespace from pathlib import Path import httpx import pytest sys.path.insert(0, str(Path(__file__).resolve().parents[1])) from app import build_apps # noqa: E402 from config import clone, load_config # noqa: E402 from services import ServiceManager # noqa: E402 class FakeVLLM(httpx.AsyncBaseTransport): """Stands in for vllm-text / vllm-ocr / vllm-embed. Endpoints: /health, /is_sleeping, /wake_up, /collective_rpc, /reset_prefix_cache, /sleep?level=, plus /v1/* which echoes what it saw. """ def __init__(self) -> None: self.calls: list[str] = [] self.services: dict[str, dict] = { "vllm-text": self._svc(), "vllm-ocr": self._svc(), "vllm-embed": self._svc(), } self.stream_chunks: list[bytes] = [] self.stream_delay: float = 0.0 # per-chunk delay for slow/SSE streams @staticmethod def _svc() -> dict: return { "reachable": True, "healthy": True, "sleeping": False, "wake_delay": 0.0, "wake_fails": 0, # number of /wake_up calls to fail (500) "sleep_fails": 0, # number of /sleep calls to fail (500) "wake_up_calls": 0, "reload_calls": 0, "reset_calls": 0, "sleep_calls": 0, "api_calls": 0, "api_delay": 0.0, } # -- helpers ----------------------------------------------------------- @staticmethod def _response(status: int = 200, *, payload: bytes = b"", content_type: str | None = None): """Build a *streamable* response. `httpx.Response(200, json=...)` would mark the body as already consumed, which the router's streaming proxy would rightly reject. """ async def gen(): if payload: yield payload headers = {"content-type": content_type} if content_type else {} return httpx.Response(status, content=gen(), headers=headers) def count(self, needle: str) -> int: return sum(1 for call in self.calls if needle in call) def last(self, needle: str) -> int: """Index of the last matching recorded call (-1 if none).""" for i in range(len(self.calls) - 1, -1, -1): if needle in self.calls[i]: return i return -1 def set_sleeping(self, key: str, sleeping: bool, level: int = 0) -> None: svc = self.services[f"vllm-{key}"] svc["sleeping"] = sleeping svc["level"] = level # -- transport --------------------------------------------------------- async def handle_async_request(self, request: httpx.Request) -> httpx.Response: host = request.url.host or "" svc = self.services.get(host) if svc is None or not svc["reachable"]: raise httpx.ConnectError(f"connection refused: {host}", request=request) path = request.url.path params = request.url.params # record method + full target (path + query) so tests can assert on # things like "level=1" target = path + (f"?{params}" if str(params) else "") self.calls.append(f"{request.method} {host}{target}") if path == "/is_sleeping": # Optional hook: report `true` for the next N polls even after a # wake-up, to model a backend whose wake is still in flight. if svc.get("hold_sleeping_polls", 0) > 0: svc["hold_sleeping_polls"] -= 1 return self._response(200, payload=b'{"is_sleeping": true}', content_type="application/json") return self._response(200, payload=json.dumps( {"is_sleeping": bool(svc["sleeping"])}).encode(), content_type="application/json") if path == "/health": # NOTE: like the real vLLM, /health answers 200 even while the # model is asleep -- readiness must be gated on /is_sleeping. if not svc["healthy"]: raise httpx.ConnectError(f"unhealthy: {host}", request=request) return self._response(200) if path == "/wake_up": svc["wake_up_calls"] += 1 if svc["wake_fails"] > 0: svc["wake_fails"] -= 1 return self._response(500, payload=b"wake failed") if svc.get("wake_flip_midway") and svc["wake_delay"]: # Model the real hazard: is_sleeping flips to false *during* # /wake_up, while reload_weights has not happened yet. A # router that dies in this window leaves a backend that looks # perfectly awake and serves garbage (E2E case 12). await asyncio.sleep(svc["wake_delay"] / 2) svc["sleeping"] = False await asyncio.sleep(svc["wake_delay"] / 2) return self._response(200) await asyncio.sleep(svc["wake_delay"]) svc["sleeping"] = False return self._response(200) if path == "/collective_rpc": svc["reload_calls"] += 1 return self._response(200) if path == "/reset_prefix_cache": svc["reset_calls"] += 1 return self._response(200) if path == "/sleep": svc["sleep_calls"] += 1 if svc["sleep_fails"] > 0: svc["sleep_fails"] -= 1 return self._response(500, payload=b"sleep failed") svc["sleeping"] = True svc["level"] = int(params.get("level", "1")) return self._response(200) if path.startswith("/v1/"): svc["api_calls"] += 1 if svc.get("api_delay"): await asyncio.sleep(svc["api_delay"]) ctype = request.headers.get("content-type", "") if "text/event-stream" in ctype or self.stream_chunks: async def gen(): for chunk in self.stream_chunks: await asyncio.sleep(self.stream_delay) yield chunk return httpx.Response(200, content=gen(), headers={"content-type": "text/event-stream"}) body = (request.content or b"").decode("utf-8", "replace") return self._response( 200, payload=json.dumps({ "echo_path": path, "echo_method": request.method, "echo_body": body, "echo_content_type": ctype, "echo_query": str(request.url.params), "service": host, }).encode(), content_type="application/json", ) return self._response(404, payload=json.dumps( {"detail": f"no route {path}"}).encode(), content_type="application/json") @pytest.fixture def backend() -> FakeVLLM: return FakeVLLM() @pytest.fixture def cfg(): config = clone(load_config()) config.state_cache_ttl = 0.0 # never trust cache -> deterministic config.wake_health_poll_s = 0.005 config.idle_enabled = False # idle is exercised via idle_tick() config.hold_sleep_s = 5.0 config.hold_offload_s = 5.0 config.hold_restart_s = 5.0 return config @pytest.fixture async def stack(cfg, backend): public_app, admin_app, manager = build_apps(cfg, transport=backend) try: yield SimpleNamespace(public=public_app, admin=admin_app, manager=manager, cfg=cfg, backend=backend) finally: await manager.proxy_client.aclose() await manager.ctrl_client.aclose() @pytest.fixture async def pub(stack): transport = httpx.ASGITransport(app=stack.public) async with httpx.AsyncClient(transport=transport, base_url="http://router.test") as client: yield client @pytest.fixture async def adm(stack): transport = httpx.ASGITransport(app=stack.admin) async with httpx.AsyncClient(transport=transport, base_url="http://router.test") as client: yield client async def raw_asgi(app, method: str, raw_path: str, *, headers=None, body: bytes = b""): """Call the ASGI app with a hand-built scope so the *raw* (percent-encoded) path reaches the router exactly as a hostile client would send it.""" scope = { "type": "http", "asgi": {"version": "3.0", "spec_version": "2.3"}, "method": method, "path": urllib.parse.unquote(raw_path), "raw_path": raw_path.encode("latin-1"), "query_string": b"", "headers": [(k.lower().encode(), v.encode()) for k, v in (headers or [])], "http_version": "1.1", "scheme": "http", "server": ("router.test", 80), "client": ("127.0.0.1", 1234), "root_path": "", } sent: list[dict] = [] incoming: list[dict] = [{"type": "http.request", "body": body, "more_body": False}] closed = asyncio.Event() async def receive(): # Behaves like a real server: hand over the body once, then wait for # the client to hang up (Starlette's disconnect listener relies on it). if incoming: return incoming.pop(0) await closed.wait() return {"type": "http.disconnect"} async def send(message): sent.append(message) await app(scope, receive, send) closed.set() status = sent[0]["status"] payload = b"".join(m.get("body", b"") for m in sent if m["type"] == "http.response.body") return status, payload