backend/tests/test_api.py

import json
import numpy as np
import pytest
from fastapi.testclient import TestClient
from pathlib import Path

import astroparse_api.main as main_mod
from astroparse_api.main import app
from astroparse_api.annotate import CompletionCache

FIXTURE = Path(__file__).parent / "fixtures" / "mowla_iyer_2024.pdf"


class FakeCorpus:
    class _DS:
        def search(self, col, q, k):
            class R:
                indices = [0]
                scores = np.array([0.5])
            return R()

        def __getitem__(self, idx):
            import datetime
            return {
                "title": ["T"],
                "abstract": ["A"],
                "authors": [["Iyer, K."]],
                "date": [datetime.date(2022, 1, 1)],
                "keywords": [["galaxy"]],
                "bibcode": ["2022ApJ...1I"],
                "arxiv_id": ["2201.00001"],
            }

    ds = _DS()

    @property
    def arxiv_ids(self) -> set:
        return {"2201.00001"}

    @property
    def bibcodes(self) -> set:
        return {"2022ApJ...1I"}

    def search(self, emb, k=1000):
        return [0], np.array([2.0])


@pytest.fixture
def client(monkeypatch, tmp_path):
    monkeypatch.setenv("ASTROPARSE_SKIP_CORPUS", "1")
    monkeypatch.setenv("ASTROPARSE_FIGURES_DIR", str(tmp_path / "figures"))
    monkeypatch.setattr(main_mod, "get_corpus", lambda: FakeCorpus())
    monkeypatch.setattr(
        main_mod,
        "embed_batch",
        lambda texts, api_key: [np.zeros(2, dtype=np.float32) for _ in texts],
    )
    main_mod.cache = CompletionCache(tmp_path / "cache.sqlite")
    with TestClient(app) as c:
        yield c


def _collect_sse(resp):
    events, cur = [], {}
    for line in resp.iter_lines():
        if line.startswith("event:"):
            cur["event"] = line.split(":", 1)[1].strip()
        elif line.startswith("data:"):
            cur["data"] = json.loads(line.split(":", 1)[1])
            events.append(cur)
            cur = {}
    return events


def test_parse_job_flow(client):
    r = client.post(
        "/api/parse",
        files={"file": ("m.pdf", FIXTURE.read_bytes(), "application/pdf")},
    )
    assert r.status_code == 202
    job_id = r.json()["jobId"]
    with client.stream("GET", f"/api/jobs/{job_id}/events") as resp:
        events = _collect_sse(resp)
    names = [e["event"] for e in events]
    assert names.count("result") == 1
    stage_events = [e for e in events if e["event"] == "stage"]
    assert {s["data"]["index"] for s in stage_events} == set(range(7))
    result = next(e for e in events if e["event"] == "result")["data"]
    assert result["paper"]["paragraphs"] and result["litByPara"]
    assert "pdfHash" in result
    assert "figures" in result
    assert "litByFig" in result


def test_annotate_cached_round_trip(client, monkeypatch):
    async def fake_stream(provider, model, prompt, key):
        for chunk in ["A fine ", "note."]:
            yield chunk

    monkeypatch.setattr(main_mod, "stream_completion", fake_stream)
    body = {
        "paragraphId": "p1",
        "paragraph": "text",
        "section": "S",
        "mode": "referee",
        "lit": [],
        "provider": "anthropic",
        "model": "claude-haiku-4-5",
        "key": "sk-x",
    }
    with client.stream("POST", "/api/annotate", json=body) as resp:
        events = _collect_sse(resp)
    done = next(e for e in events if e["event"] == "done")["data"]
    assert done["text"] == "A fine note." and done["cached"] is False

    with client.stream("POST", "/api/annotate", json=body) as resp:  # second call: cache hit
        events2 = _collect_sse(resp)
    done2 = next(e for e in events2 if e["event"] == "done")["data"]
    assert done2["cached"] is True and done2["text"] == "A fine note."


def test_annotate_missing_key(client):
    body = {
        "paragraphId": "p1",
        "paragraph": "t",
        "section": "S",
        "mode": "referee",
        "lit": [],
        "provider": "anthropic",
        "model": "m",
        "key": "",
    }
    with client.stream("POST", "/api/annotate", json=body) as resp:
        events = _collect_sse(resp)
    assert events[-1]["event"] == "error"
    assert "no API key set" in events[-1]["data"]["message"]


def test_figures_endpoint_serves_png(client, tmp_path, monkeypatch):
    """After a parse, the figure PNG can be fetched from the figures endpoint."""
    # Run a parse to generate figures
    r = client.post(
        "/api/parse",
        files={"file": ("m.pdf", FIXTURE.read_bytes(), "application/pdf")},
    )
    job_id = r.json()["jobId"]
    with client.stream("GET", f"/api/jobs/{job_id}/events") as resp:
        events = _collect_sse(resp)
    result = next(e for e in events if e["event"] == "result")["data"]
    pdf_hash = result["pdfHash"]
    figures = result["figures"]

    # Find a figure with an image
    with_image = [f for f in figures if f["hasImage"]]
    assert with_image, "Expected at least one figure with an image"
    fig_id = with_image[0]["id"]

    resp = client.get(f"/api/figures/{pdf_hash}/{fig_id}.png")
    assert resp.status_code == 200
    assert resp.headers["content-type"] == "image/png"
    assert len(resp.content) > 5_000


def test_figures_endpoint_invalid_hash(client):
    """An invalid (non-hex-64) hash returns 404."""
    r = client.get("/api/figures/invalid_hash/f1.png")
    assert r.status_code == 404


def test_figures_endpoint_invalid_fig_id(client):
    """An invalid fig_id like 'fx1' returns 404."""
    good_hash = "a" * 64
    r = client.get(f"/api/figures/{good_hash}/fx1.png")
    assert r.status_code == 404


def test_figures_endpoint_traversal_attempt(client):
    """Traversal payloads in either path segment are rejected by validation."""
    good_hash = "a" * 64
    # traversal in the hash segment (encoded so the path reaches the route)
    r = client.get(f"/api/figures/..%2f..%2f{good_hash[:58]}/f1.png")
    assert r.status_code == 404
    # traversal in the fig_id segment
    r = client.get(f"/api/figures/{good_hash}/..%2f..%2fsecret.png")
    assert r.status_code == 404
    # valid hash + valid fig_id but no file on disk → also 404 (no info leak)
    r = client.get(f"/api/figures/{good_hash}/f1.png")
    assert r.status_code == 404


def test_jobs_stage_count():
    """STAGES list must have exactly 7 entries."""
    from astroparse_api.jobs import STAGES
    assert len(STAGES) == 7
    assert STAGES[6] == "Tipping in the plates"


# ---------------------------------------------------------------------------
# Feature: plaintext/markdown manuscript input
# ---------------------------------------------------------------------------

_MD_UPLOAD = (
    "# Test Manuscript\n\n"
    "## Introduction\n\n"
    + ("Galaxy formation and evolution in the early universe is a key topic. " * 20)
    + "\n\n"
    + ("Star formation rates and quenching mechanisms drive cosmic evolution. " * 20)
    + "\n\n"
    "## References\n\n"
    "[1] Someone et al. 2020\n"
).encode()


def test_markdown_upload_all_stages_emitted(client):
    """Uploading a .md file triggers all 7 stages and returns figures=[]."""
    r = client.post(
        "/api/parse",
        files={"file": ("notes.md", _MD_UPLOAD, "text/markdown")},
    )
    assert r.status_code == 202
    job_id = r.json()["jobId"]
    with client.stream("GET", f"/api/jobs/{job_id}/events") as resp:
        events = _collect_sse(resp)
    # All 7 stages must be emitted
    stage_events = [e for e in events if e["event"] == "stage"]
    assert {s["data"]["index"] for s in stage_events} == set(range(7)), (
        f"Missing stages: got indices {[s['data']['index'] for s in stage_events]}"
    )
    # Result must have paragraphs, figures=[], litByPara set
    result = next(e for e in events if e["event"] == "result")["data"]
    assert result["paper"]["paragraphs"], "Expected paragraphs in result"
    assert result["figures"] == [], f"Expected figures=[], got {result['figures']}"
    assert result["litByPara"], "Expected litByPara non-empty"
    assert result["litByFig"] == {}, f"Expected litByFig={{}}, got {result['litByFig']}"


def test_annotate_network_error_sanitized(client, monkeypatch):
    """Non-ProviderError failures (e.g. httpx ConnectError) emit a sanitized
    error event — never the raw exception, which can carry a key-bearing URL."""
    async def broken_stream(provider, model, prompt, key):
        raise RuntimeError(f"https://example/?key={key}")
        yield  # pragma: no cover — makes this an async generator
    monkeypatch.setattr(main_mod, "stream_completion", broken_stream)
    body = {
        "paragraphId": "p1", "paragraph": "t", "section": "S", "mode": "referee",
        "lit": [], "provider": "gemini", "model": "m", "key": "SECRETKEY",
    }
    with client.stream("POST", "/api/annotate", json=body) as resp:
        events = _collect_sse(resp)
    assert events[-1]["event"] == "error"
    msg = events[-1]["data"]["message"]
    assert "SECRETKEY" not in msg and "could not reach gemini" in msg