diff --git a/tests/test_app.py b/tests/test_app.py index dd3c36f..e6131b4 100644 --- a/tests/test_app.py +++ b/tests/test_app.py @@ -1,10 +1,12 @@ import io +import json from pathlib import Path import pytest from httpx import ASGITransport, AsyncClient -from unredact.app import app +import unredact.app +from unredact.app import app, PAGE_SIZE, _solve_results @pytest.fixture @@ -261,3 +263,125 @@ async def test_spot_returns_analysis(pdf_bytes: bytes): assert "segments" in data["analysis"] assert "gap" in data["analysis"] assert "line" in data["analysis"] + + +@pytest.mark.anyio +async def test_solve_paginates_results(): + """Solve should stream at most PAGE_SIZE results and buffer the rest.""" + from unredact.app import PAGE_SIZE + + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url="http://test") as client: + resp = await client.get("/api/fonts") + fonts = resp.json()["fonts"] + font = next(f for f in fonts if f["available"]) + + # Use word mode with full vocab and wide tolerance to get many results + resp = await client.post("/api/solve", json={ + "font_id": font["id"], + "font_size": 40, + "gap_width_px": 80.0, + "tolerance_px": 10.0, + "left_context": "", + "right_context": "", + "hints": {"charset": "lowercase"}, + "mode": "word", + "vocab_size": 0, + }) + assert resp.status_code == 200 + + # Parse SSE events + events = [] + for line in resp.text.split("\n"): + if line.startswith("data: "): + try: + events.append(json.loads(line[6:])) + except (json.JSONDecodeError, ValueError): + pass + + matches = [e for e in events if e.get("status") == "match"] + done_events = [e for e in events if e.get("status") == "done"] + page_events = [e for e in events if e.get("status") == "page_complete"] + + assert len(done_events) == 1 + total = done_events[0]["total_found"] + + if total > PAGE_SIZE: + assert len(matches) == PAGE_SIZE + assert len(page_events) == 1 + assert page_events[0]["sent"] == PAGE_SIZE + else: + assert len(matches) == total + assert len(page_events) == 0 + + +def _parse_sse_events(body: str) -> list[dict]: + """Parse SSE body text into a list of JSON event dicts.""" + events = [] + for line in body.splitlines(): + line = line.strip() + if line.startswith("data: "): + try: + events.append(json.loads(line[6:])) + except (json.JSONDecodeError, ValueError): + continue + return events + + +@pytest.mark.anyio +async def test_solve_pagination_caps_at_page_size(monkeypatch): + """POST /api/solve should cap streamed match events at PAGE_SIZE and send page_complete.""" + # Use a small PAGE_SIZE so we don't need thousands of results + small_page = 5 + monkeypatch.setattr(unredact.app, "PAGE_SIZE", small_page) + + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url="http://test") as client: + resp = await client.get("/api/fonts") + fonts = resp.json()["fonts"] + font = next(f for f in fonts if f["available"]) + + # Word mode with small vocab but wide tolerance to generate many matches (>> 5) + resp = await client.post("/api/solve", json={ + "font_id": font["id"], + "font_size": 40, + "gap_width_px": 100.0, + "tolerance_px": 20.0, + "left_context": "", + "right_context": "", + "hints": {"charset": "lowercase"}, + "mode": "word", + "vocab_size": 200, + }) + assert resp.status_code == 200 + assert "text/event-stream" in resp.headers.get("content-type", "") + + solve_id = resp.headers["x-solve-id"] + events = _parse_sse_events(resp.text) + + # Separate event types + match_events = [e for e in events if e.get("status") == "match"] + page_complete_events = [e for e in events if e.get("status") == "page_complete"] + done_events = [e for e in events if e.get("status") == "done"] + + # Exactly PAGE_SIZE match events should be streamed + assert len(match_events) == small_page, ( + f"Expected {small_page} match events, got {len(match_events)}" + ) + + # A page_complete event should be present + assert len(page_complete_events) == 1 + pc = page_complete_events[0] + assert pc["sent"] == small_page + assert pc["solve_id"] == solve_id + + # The done event should still be present with total_found > PAGE_SIZE + assert len(done_events) == 1 + total = done_events[0]["total_found"] + assert total > small_page, ( + f"Expected total_found > {small_page}, got {total}" + ) + + # The buffer should contain ALL results (including the ones streamed) + assert solve_id in _solve_results + assert len(_solve_results[solve_id]) == total diff --git a/unredact/app.py b/unredact/app.py index ff53e60..48d424b 100644 --- a/unredact/app.py +++ b/unredact/app.py @@ -370,10 +370,29 @@ async def solve(req: SolveRequest): _active_solves[solve_id] = False async def event_generator(): + _solve_results.clear() + _solve_results[solve_id] = [] + try: found_texts = set() + sent_count = 0 + page_complete_sent = False charset_name = req.hints.get("charset", "lowercase") + def _record(d): + """Buffer result and return True if it should be yielded as SSE.""" + nonlocal sent_count, page_complete_sent + _solve_results[solve_id].append(d) + sent_count += 1 + return sent_count <= PAGE_SIZE + + def _page_complete_event(): + nonlocal page_complete_sent + if not page_complete_sent: + page_complete_sent = True + return json.dumps({"status": "page_complete", "sent": PAGE_SIZE, "solve_id": solve_id}) + return None + # Name mode: single-word associate names if req.mode == "name" and not _active_solves.get(solve_id): name_results = solve_name_dictionary( @@ -389,13 +408,19 @@ async def solve(req: SolveRequest): if r.text in found_texts: continue found_texts.add(r.text) - yield json.dumps({ + d = { "status": "match", "text": r.text, "width_px": round(r.width, 2), "error_px": round(r.error, 2), "source": "names", - }) + } + if _record(d): + yield json.dumps(d) + else: + evt = _page_complete_event() + if evt: + yield evt # Full name mode: first x last Cartesian product + variants if req.mode == "full_name" and not _active_solves.get(solve_id): @@ -412,13 +437,19 @@ async def solve(req: SolveRequest): if r.text in found_texts: continue found_texts.add(r.text) - yield json.dumps({ + d = { "status": "match", "text": r.text, "width_px": round(r.width, 2), "error_px": round(r.error, 2), "source": "names", - }) + } + if _record(d): + yield json.dumps(d) + else: + evt = _page_complete_event() + if evt: + yield evt # Email mode if req.mode == "email" and not _active_solves.get(solve_id): @@ -432,13 +463,19 @@ async def solve(req: SolveRequest): if _active_solves.get(solve_id): break found_texts.add(r.text) - yield json.dumps({ + d = { "status": "match", "text": r.text, "width_px": round(r.width, 2), "error_px": round(r.error, 2), "source": "emails", - }) + } + if _record(d): + yield json.dumps(d) + else: + evt = _page_complete_event() + if evt: + yield evt # Word mode: English nouns + adjective/noun phrases if req.mode == "word" and not _active_solves.get(solve_id): @@ -456,13 +493,19 @@ async def solve(req: SolveRequest): if r.text in found_texts: continue found_texts.add(r.text) - yield json.dumps({ + d = { "status": "match", "text": r.text, "width_px": round(r.width, 2), "error_px": round(r.error, 2), "source": "words", - }) + } + if _record(d): + yield json.dumps(d) + else: + evt = _page_complete_event() + if evt: + yield evt # Enumerate mode: Rust backend if req.mode == "enumerate" and not _active_solves.get(solve_id): @@ -502,13 +545,19 @@ async def solve(req: SolveRequest): if text in found_texts: continue found_texts.add(text) - yield json.dumps({ + d = { "status": "match", "text": text, "width_px": round(r["width"], 2), "error_px": round(r["error"], 2), "source": "enumerate", - }) + } + if _record(d): + yield json.dumps(d) + else: + evt = _page_complete_event() + if evt: + yield evt yield json.dumps({ "status": "done",