94 lines
3.4 KiB
Python
94 lines
3.4 KiB
Python
"""Phase 64 E2E helper — a delay-injecting reverse proxy in front of the
|
||
mock LLM.
|
||
|
||
The mock LLM (``mock_llm.py``) answers instantly: a real archive scan of
|
||
dozens of files finishes in well under a second and never outlives the
|
||
UI's 2 s status poll — the phase-64 live file labels ("Processing…
|
||
<file>", "Importing <file>", "Syncing… <file>"), the at-202 toast →
|
||
navigate-away contract, and the mid-scan reload re-attach would all be
|
||
races against the mock. This proxy sits between the app under test and
|
||
the mock LLM and sleeps ``SLOW_LLM_DELAY_S`` seconds (default 0.15)
|
||
before forwarding each request, so a scan's duration is deterministic
|
||
(≈ the run's number of LLM requests × the delay — for an N-file
|
||
archive/sync that is N + 3: the ``check_models`` embed + chat probe,
|
||
one embed per file, and the change-gated overview chat). Both the
|
||
tests' ~100 ms status polling (the deterministic layer) and the UI's
|
||
2 s poll (the UI layer) then observe the running state, the current
|
||
file, and the counts reliably.
|
||
|
||
Everything else is byte-transparent: method, path, query, headers, and
|
||
body are forwarded verbatim; the upstream response's status and body
|
||
come back as-is (``content-encoding`` / ``content-length`` are dropped
|
||
— httpx has already decoded the body and starlette recomputes the
|
||
length). The mock's responses are all finite (its SSE streams end with
|
||
``[DONE]``), so the proxy reads each body to completion before
|
||
answering.
|
||
|
||
Run it the conftest way (a suite's module ``app_server`` fixture spawns
|
||
it as a subprocess):
|
||
|
||
uv run python -m uvicorn tests.e2e.slow_llm:app --port 8902
|
||
|
||
with ``E2E_MOCK_PORT`` (upstream, default 8901) and ``SLOW_LLM_DELAY_S``
|
||
(delay, default 0.15) in its environment.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import os
|
||
|
||
import httpx
|
||
from fastapi import FastAPI, Request, Response
|
||
|
||
#: The mock LLM this proxy forwards to (the conftest's MOCK_PORT).
|
||
MOCK_PORT = int(os.environ.get("E2E_MOCK_PORT", "8901"))
|
||
UPSTREAM = f"http://127.0.0.1:{MOCK_PORT}"
|
||
|
||
#: Per-request delay in seconds — each suite's proxy fixture picks its
|
||
#: own (passed through the subprocess env).
|
||
DELAY_S = float(os.environ.get("SLOW_LLM_DELAY_S", "0.15"))
|
||
|
||
app = FastAPI()
|
||
_client: httpx.AsyncClient | None = None
|
||
|
||
|
||
def _get_client() -> httpx.AsyncClient:
|
||
global _client
|
||
if _client is None:
|
||
_client = httpx.AsyncClient(base_url=UPSTREAM, timeout=60.0)
|
||
return _client
|
||
|
||
|
||
@app.on_event("shutdown")
|
||
async def _close_client() -> None:
|
||
global _client
|
||
if _client is not None:
|
||
await _client.aclose()
|
||
_client = None
|
||
|
||
|
||
@app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH"])
|
||
async def proxy(path: str, request: Request) -> Response:
|
||
"""Sleep ``DELAY_S``, then forward the request to the mock LLM."""
|
||
await asyncio.sleep(DELAY_S)
|
||
body = await request.body()
|
||
headers = {
|
||
k: v for k, v in request.headers.items() if k.lower() not in ("host", "content-length")
|
||
}
|
||
upstream = await _get_client().request(
|
||
request.method,
|
||
f"/{path}",
|
||
content=body,
|
||
headers=headers,
|
||
params=dict(request.query_params),
|
||
)
|
||
resp_headers = {
|
||
k: v
|
||
for k, v in upstream.headers.items()
|
||
if k.lower()
|
||
not in ("content-length", "content-encoding", "transfer-encoding", "connection")
|
||
}
|
||
return Response(
|
||
content=upstream.content, status_code=upstream.status_code, headers=resp_headers
|
||
)
|