mirror of
https://github.com/PR0M3TH3AN/unredact.git
synced 2026-10-27 08:40:29 +00:00
Alex-GilbertandClaude Opus 4.6
672bdc3223
- Batch candidates in groups of 50 (was 200) to avoid output token
truncation that caused all-zero scores
- Stream scored results via SSE as each batch completes instead of
waiting for all batches to finish
- Include solve_id in done event so validate works when results fit
in a single page
- Rewrite prompt to judge sentence coherence ("does it make sense")
rather than word predictability ("is it expected")
- Frontend inserts results in score-sorted order as they arrive
Co-Authored-By: Claude Opus 4.6 <[email protected]>
510 lines
18 KiB
Python
510 lines
18 KiB
Python
import io
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
from httpx import ASGITransport, AsyncClient
|
|
|
|
import unredact.app
|
|
from unredact.app import app, PAGE_SIZE, _solve_results
|
|
|
|
|
|
@pytest.fixture
|
|
def pdf_bytes(sample_pdf: Path) -> bytes:
|
|
return sample_pdf.read_bytes()
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_upload_pdf(pdf_bytes: bytes):
|
|
transport = ASGITransport(app=app)
|
|
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
|
resp = await client.post(
|
|
"/api/upload",
|
|
files={"file": ("test.pdf", pdf_bytes, "application/pdf")},
|
|
)
|
|
assert resp.status_code == 200
|
|
data = resp.json()
|
|
assert "doc_id" in data
|
|
assert data["page_count"] > 0
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_get_page_original(pdf_bytes: bytes):
|
|
transport = ASGITransport(app=app)
|
|
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
|
resp = await client.post(
|
|
"/api/upload",
|
|
files={"file": ("test.pdf", pdf_bytes, "application/pdf")},
|
|
)
|
|
doc_id = resp.json()["doc_id"]
|
|
|
|
resp = await client.get(f"/api/doc/{doc_id}/page/1/original")
|
|
assert resp.status_code == 200
|
|
assert resp.headers["content-type"] == "image/png"
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_ocr_endpoint_streams_results(pdf_bytes: bytes):
|
|
"""GET /api/doc/{id}/ocr should stream OCR progress and cache results."""
|
|
transport = ASGITransport(app=app)
|
|
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
|
resp = await client.post(
|
|
"/api/upload",
|
|
files={"file": ("test.pdf", pdf_bytes, "application/pdf")},
|
|
)
|
|
doc_id = resp.json()["doc_id"]
|
|
|
|
resp = await client.get(f"/api/doc/{doc_id}/ocr")
|
|
assert resp.status_code == 200
|
|
assert "text/event-stream" in resp.headers.get("content-type", "")
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_get_page_data_returns_redactions(pdf_bytes: bytes):
|
|
"""Page data should return redaction bboxes, not OCR lines."""
|
|
transport = ASGITransport(app=app)
|
|
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
|
resp = await client.post(
|
|
"/api/upload",
|
|
files={"file": ("test.pdf", pdf_bytes, "application/pdf")},
|
|
)
|
|
doc_id = resp.json()["doc_id"]
|
|
|
|
resp = await client.get(f"/api/doc/{doc_id}/page/1/data")
|
|
assert resp.status_code == 200
|
|
data = resp.json()
|
|
assert "redactions" in data
|
|
assert isinstance(data["redactions"], list)
|
|
if data["redactions"]:
|
|
r = data["redactions"][0]
|
|
assert "id" in r
|
|
assert "x" in r
|
|
assert "y" in r
|
|
assert "w" in r
|
|
assert "h" in r
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_redaction_analyze(pdf_bytes: bytes):
|
|
"""POST /api/redaction/analyze should OCR the line around a redaction."""
|
|
transport = ASGITransport(app=app)
|
|
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
|
resp = await client.post(
|
|
"/api/upload",
|
|
files={"file": ("test.pdf", pdf_bytes, "application/pdf")},
|
|
)
|
|
doc_id = resp.json()["doc_id"]
|
|
|
|
resp = await client.get(f"/api/doc/{doc_id}/page/1/data")
|
|
redactions = resp.json()["redactions"]
|
|
if not redactions:
|
|
pytest.skip("No redactions detected in test PDF")
|
|
|
|
r = redactions[0]
|
|
|
|
resp = await client.post("/api/redaction/analyze", json={
|
|
"doc_id": doc_id,
|
|
"page": 1,
|
|
"redaction": {"x": r["x"], "y": r["y"], "w": r["w"], "h": r["h"]},
|
|
})
|
|
assert resp.status_code == 200
|
|
data = resp.json()
|
|
assert "segments" in data
|
|
assert "gap" in data
|
|
assert "font" in data
|
|
assert data["gap"]["w"] > 0
|
|
assert data["font"]["name"]
|
|
assert data["font"]["size"] > 0
|
|
assert "offset_x" in data
|
|
assert "offset_y" in data
|
|
assert isinstance(data["offset_x"], (int, float))
|
|
assert isinstance(data["offset_y"], (int, float))
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_list_fonts():
|
|
"""GET /api/fonts should return list of candidate fonts."""
|
|
transport = ASGITransport(app=app)
|
|
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
|
resp = await client.get("/api/fonts")
|
|
assert resp.status_code == 200
|
|
data = resp.json()
|
|
assert "fonts" in data
|
|
assert len(data["fonts"]) > 0
|
|
font = data["fonts"][0]
|
|
assert "name" in font
|
|
assert "id" in font
|
|
assert "available" in font
|
|
# At least one font should be available on the system
|
|
assert any(f["available"] for f in data["fonts"])
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_get_font_ttf():
|
|
"""GET /api/font/{id} should serve a TTF file."""
|
|
transport = ASGITransport(app=app)
|
|
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
|
# Get list first to find an available font
|
|
resp = await client.get("/api/fonts")
|
|
fonts = resp.json()["fonts"]
|
|
available = [f for f in fonts if f["available"]]
|
|
assert available, "No fonts available on system"
|
|
|
|
font_id = available[0]["id"]
|
|
resp = await client.get(f"/api/font/{font_id}")
|
|
assert resp.status_code == 200
|
|
assert resp.headers["content-type"] == "font/ttf"
|
|
# TTF files start with specific bytes or have substantial size
|
|
assert len(resp.content) > 1000
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_get_font_not_found():
|
|
"""GET /api/font/{id} with bad id should 404."""
|
|
transport = ASGITransport(app=app)
|
|
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
|
resp = await client.get("/api/font/nonexistent-font")
|
|
assert resp.status_code == 404
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_solve_endpoint_enumerate():
|
|
"""POST /api/solve should stream SSE results."""
|
|
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"])
|
|
|
|
resp = await client.post("/api/solve", json={
|
|
"font_id": font["id"],
|
|
"font_size": 40,
|
|
"gap_width_px": 50.0,
|
|
"tolerance_px": 5.0,
|
|
"left_context": "",
|
|
"right_context": "",
|
|
"hints": {
|
|
"charset": "lowercase",
|
|
},
|
|
"mode": "enumerate",
|
|
})
|
|
assert resp.status_code == 200
|
|
assert "text/event-stream" in resp.headers.get("content-type", "")
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_solve_endpoint_name_mode():
|
|
"""POST /api/solve with mode=name should work."""
|
|
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"])
|
|
|
|
resp = await client.post("/api/solve", json={
|
|
"font_id": font["id"],
|
|
"font_size": 40,
|
|
"gap_width_px": 50.0,
|
|
"tolerance_px": 5.0,
|
|
"left_context": "",
|
|
"right_context": "",
|
|
"hints": {
|
|
"charset": "lowercase",
|
|
},
|
|
"mode": "name",
|
|
})
|
|
assert resp.status_code == 200
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_get_associates():
|
|
"""GET /api/associates should return the associates lookup data."""
|
|
transport = ASGITransport(app=app)
|
|
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
|
resp = await client.get("/api/associates")
|
|
assert resp.status_code == 200
|
|
data = resp.json()
|
|
assert "names" in data
|
|
assert "persons" in data
|
|
assert isinstance(data["names"], dict)
|
|
assert isinstance(data["persons"], dict)
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_spot_returns_analysis(pdf_bytes: bytes):
|
|
"""POST /spot should return analysis data when OCR is cached."""
|
|
transport = ASGITransport(app=app)
|
|
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
|
resp = await client.post(
|
|
"/api/upload",
|
|
files={"file": ("test.pdf", pdf_bytes, "application/pdf")},
|
|
)
|
|
doc_id = resp.json()["doc_id"]
|
|
|
|
# Run OCR first (required for spot analysis)
|
|
resp = await client.get(f"/api/doc/{doc_id}/ocr")
|
|
assert resp.status_code == 200
|
|
|
|
# Try to spot a redaction at an arbitrary position
|
|
resp = await client.post(
|
|
f"/api/doc/{doc_id}/page/1/spot",
|
|
json={"x": 300, "y": 300},
|
|
)
|
|
if resp.status_code == 404:
|
|
pytest.skip("No redaction found at test coordinates")
|
|
|
|
data = resp.json()
|
|
assert "id" in data
|
|
assert "x" in data
|
|
assert "analysis" in data
|
|
# analysis should be present (or null if no OCR line matched)
|
|
if data["analysis"] is not None:
|
|
assert "font" in data["analysis"]
|
|
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
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_solve_pagination_endpoint():
|
|
"""GET /api/solve/{id}/results should return paginated buffered results."""
|
|
fake_id = "test123"
|
|
_solve_results[fake_id] = [
|
|
{"text": f"word{i}", "width_px": 50.0, "error_px": 0.1, "source": "words"}
|
|
for i in range(500)
|
|
]
|
|
|
|
transport = ASGITransport(app=app)
|
|
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
|
# First page
|
|
resp = await client.get(f"/api/solve/{fake_id}/results?offset=0&limit=200")
|
|
assert resp.status_code == 200
|
|
data = resp.json()
|
|
assert len(data["results"]) == 200
|
|
assert data["total"] == 500
|
|
assert data["offset"] == 0
|
|
assert data["limit"] == 200
|
|
assert data["complete"] is True
|
|
|
|
# Second page
|
|
resp = await client.get(f"/api/solve/{fake_id}/results?offset=200&limit=200")
|
|
data = resp.json()
|
|
assert len(data["results"]) == 200
|
|
assert data["results"][0]["text"] == "word200"
|
|
|
|
# Third page (partial)
|
|
resp = await client.get(f"/api/solve/{fake_id}/results?offset=400&limit=200")
|
|
data = resp.json()
|
|
assert len(data["results"]) == 100
|
|
|
|
# Unknown solve_id
|
|
resp = await client.get("/api/solve/nonexistent/results")
|
|
assert resp.status_code == 404
|
|
|
|
_solve_results.pop(fake_id, None)
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_cancel_clears_buffer():
|
|
"""DELETE /api/solve/{id} should also clean up the results buffer."""
|
|
from unredact.app import _active_solves
|
|
|
|
fake_id = "cancel_test"
|
|
_active_solves[fake_id] = False
|
|
_solve_results[fake_id] = [{"text": "test"}]
|
|
|
|
transport = ASGITransport(app=app)
|
|
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
|
resp = await client.delete(f"/api/solve/{fake_id}")
|
|
assert resp.status_code == 200
|
|
|
|
assert fake_id not in _solve_results
|
|
_active_solves.pop(fake_id, None)
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_validate_endpoint(monkeypatch):
|
|
"""POST /api/solve/{id}/validate should stream SSE with LLM-scored results."""
|
|
from unredact.pipeline import llm_validate
|
|
|
|
fake_id = "validate_test"
|
|
_solve_results[fake_id] = [
|
|
{"text": "Smith", "width_px": 50.0, "error_px": 0.1, "source": "names"},
|
|
{"text": "house", "width_px": 50.0, "error_px": 0.2, "source": "words"},
|
|
{"text": "running", "width_px": 50.0, "error_px": 0.3, "source": "words"},
|
|
]
|
|
|
|
async def mock_score_batch(client, left, right, candidates):
|
|
scores = {"Smith": 95, "house": 10, "running": 5}
|
|
return [scores.get(c, 0) for c in candidates]
|
|
|
|
monkeypatch.setattr(llm_validate, "_score_batch", mock_score_batch)
|
|
|
|
# Mock _get_client since the endpoint imports it directly
|
|
from unredact.pipeline import llm_detect
|
|
monkeypatch.setattr(llm_detect, "_get_client", lambda: None)
|
|
|
|
transport = ASGITransport(app=app)
|
|
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
|
resp = await client.post(f"/api/solve/{fake_id}/validate", json={
|
|
"left_context": "Dear Mr.",
|
|
"right_context": ", we are writing",
|
|
})
|
|
assert resp.status_code == 200
|
|
|
|
# Parse SSE events from response
|
|
events = []
|
|
for line in resp.text.split("\n"):
|
|
if line.startswith("data: "):
|
|
events.append(json.loads(line[6:]))
|
|
|
|
batch_events = [e for e in events if e.get("status") == "batch_done"]
|
|
assert len(batch_events) == 1 # 3 items fits in one batch
|
|
data = batch_events[0]
|
|
assert len(data["results"]) == 3
|
|
# Results within a batch are unsorted (frontend sorts on insert)
|
|
scores = {r["text"]: r["llm_score"] for r in data["results"]}
|
|
assert scores["Smith"] == 95
|
|
assert scores["house"] == 10
|
|
assert scores["running"] == 5
|
|
|
|
done_event = [e for e in events if e.get("status") == "done"]
|
|
assert len(done_event) == 1
|
|
assert done_event[0]["total"] == 3
|
|
|
|
_solve_results.pop(fake_id, None)
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_validate_unknown_solve_id():
|
|
"""POST /api/solve/{id}/validate with unknown id should 404."""
|
|
transport = ASGITransport(app=app)
|
|
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
|
resp = await client.post("/api/solve/nonexistent/validate", json={
|
|
"left_context": "test",
|
|
"right_context": "test",
|
|
})
|
|
assert resp.status_code == 404
|