test_figures_e2e.py

"""E2E test of the in-cycle figure stage against a local Postgres.

Vision calls (classify + mine) are stubbed at figures.py's public seam;
everything else is real code: harvest on a real generated PDF (embedded
raster + 'Figure 1' caption), figure_properties_to_rows (canonicalize +
plausibility), the pg_mirror figure backend (upsert with crop bytes,
origin-aware dedup, done-tracking), orchestrator metrics, run bookkeeping,
the publish-guard downgrade, and origin-scoped promote SQL.

Run with the same env as test_e2e.py (DB_* pointing at a scratch Postgres).
"""
import os
import sys

os.environ.setdefault("AGENT_WORK_DIR", "/tmp/agent_fig_test")
os.environ["GEMINI_API_KEY"] = "test-key-not-used"
os.environ["AGENT_USE_NTRS"] = "0"      # tests stub the lanes; never touch the network
# No network in tests: the live sources added Sep 2026 (Semantic Scholar
# bulk, OpenAlex topic walk, datasheet seed crawl) are off, and the query
# grid is not seeded so the frontier is the 8 builtin intents.
os.environ["AGENT_USE_S2"] = "0"
os.environ["AGENT_USE_OPENALEX_TOPICS"] = "0"
os.environ["AGENT_USE_DATASHEETS"] = "0"
os.environ["AGENT_SEED_QUERY_GRID"] = "0"

import fitz
import pdf_crawler
import extraction
import batch_ingest
import figures
import pg_mirror
from agent import agdb, config as C, orchestrator

C.GEMINI_API_KEY = "test-key-not-used"

QUOTE = "The tensile modulus of PEEK-CF30 was measured as 24.5 GPa at 23 C."
CAPTION = "Figure 1. Tensile strength of PEEK-CF30 versus temperature."


def make_plot_png(w: int = 400, h: int = 300) -> bytes:
    """A plot-ish raster: white ground, axes, a rising curve."""
    doc = fitz.open()
    page = doc.new_page(width=w, height=h)
    page.draw_rect(fitz.Rect(0, 0, w, h), color=None, fill=(1, 1, 1))
    page.draw_line(fitz.Point(40, h - 40), fitz.Point(w - 20, h - 40), width=2)
    page.draw_line(fitz.Point(40, 20), fitz.Point(40, h - 40), width=2)
    pts = [fitz.Point(40 + i * (w - 70) / 10, h - 60 - i * (h - 100) / 12)
           for i in range(11)]
    for a, b in zip(pts, pts[1:]):
        page.draw_line(a, b, color=(0, 0, 1), width=2)
    png = page.get_pixmap().tobytes("png")
    doc.close()
    return png


def make_pdf() -> bytes:
    doc = fitz.open()
    page = doc.new_page()  # 612 x 792 pt
    page.insert_text((72, 90), "PEEK CF30 composite datasheet")
    page.insert_text((72, 115), QUOTE)
    # Embedded raster big enough for every junk filter
    rect = fitz.Rect(72, 150, 392, 390)   # 320 x 240 pt ≈ 9.9% of the page
    page.insert_image(rect, stream=make_plot_png())
    page.insert_text((72, 408), CAPTION)
    data = doc.tobytes()
    doc.close()
    return data


FAKE_PDF = make_pdf()

# --- discovery / download / text-extraction fakes (as in test_e2e.py) -------


def fake_search_openalex(query, limit):
    yield pdf_crawler.Candidate(
        title="Figure mining test PEEK CF30 thermoplastic composite",
        pdf_url="https://example.org/peek_cf30_fig.pdf", source="openalex",
        query=query, doi="10.9999/aim.test.0002", year="2026",
        abstract="thermoplastic composite tensile strength carbon fiber PEEK plot")


def fake_search_arxiv(query, limit):
    return iter(())


def fake_download(cand, pdf_dir, state):
    import hashlib
    if cand.pdf_url in state.seen_urls:
        return None
    state.seen_urls.add(cand.pdf_url)
    sha = hashlib.sha256(FAKE_PDF).hexdigest()
    if sha in state.seen_hashes:
        return None
    state.seen_hashes.add(sha)
    fname = f"{cand.source}_{pdf_crawler.slugify(cand.title)}_{sha[:8]}.pdf"
    (pdf_dir / fname).write_bytes(FAKE_PDF)
    return {"filename": fname, "title": cand.title, "doi": cand.doi,
            "url": cand.pdf_url, "year": cand.year, "source": cand.source,
            "sha256": sha, "query": cand.query, "bytes": len(FAKE_PDF)}


def fake_extract(pdf_bytes, filename, api_key):
    mat = extraction.Material(
        material_name="Polyetheretherketone 30% Carbon Fiber",
        material_abbreviation="PEEK-CF30", material_class="Composite",
        matrix="PEEK", fiber="Carbon", fiber_volume_fraction="30%",
        properties=[
            extraction.Property(
                section="Mechanical", property_name="Tensile Modulus",
                value_raw="24.5", value_num=24.5, unit="GPa",
                test_condition="23 C", source_quote=QUOTE, page=1),
        ])
    return extraction.Extraction(materials=[mat], doc_status="ok")


# --- vision fakes at figures.py's public seam --------------------------------

CALLS = {"classify": 0, "mine": 0}


def fake_classify(figs, text_materials, api_key, *, model="m", stats=None):
    CALLS["classify"] += 1
    if stats is not None:
        stats.classify_calls += 1
    for f in figs:
        f.figure_kind = "property_plot"
        f.material_guess = "PEEK-CF30"
        f.mining_status = ("not_mined" if f.figure_kind in figures.MINEABLE_KINDS
                           else "skipped_kind")
    return figs


def fake_mine(fig, text_materials, api_key, *, model="m", stats=None):
    CALLS["mine"] += 1
    if stats is not None:
        stats.mining_calls += 1
    fig.mining_status = "mined"
    fig.n_values = 1
    return figures.MinedFigure(
        figure=fig, x_label="Temperature", x_unit="degC",
        y_label="Tensile Strength", y_unit="MPa",
        values=[("PEEK-CF30", "PEEK-CF30 curve", extraction.Property(
            section="Mechanical", property_name="Tensile Strength",
            value_raw="85", value_num=85.0, unit="MPa",
            test_condition="150 C", source_quote="", page=fig.page,
            comments="peak of curve"))])


# --- monkeypatch ------------------------------------------------------------
pdf_crawler.search_openalex = fake_search_openalex
pdf_crawler.search_arxiv = fake_search_arxiv
pdf_crawler.download_pdf = fake_download
batch_ingest.extract_from_pdf = fake_extract
figures.classify_figures = fake_classify
figures.mine_figure = fake_mine

# --- run one cycle ----------------------------------------------------------
orchestrator.bootstrap()
# This suite exercises the mining path; mining is off by default since the
# 28 Sep 2026 cost defaults (Run Control can switch it on, as here).
_c = agdb.connect()
agdb.set_config(_c, {"figure_mining": True})
_c.close()
m1 = orchestrator.run_cycle(trigger="test-fig")
print("cycle metrics:", m1)

fails = []


def check(name, cond):
    print(("PASS " if cond else "FAIL ") + name)
    if not cond:
        fails.append(name)


conn = agdb.connect()
try:
    with conn.cursor() as cur:
        cur.execute("SELECT figure_id, caption, figure_kind, mining_status, "
                    "octet_length(image_bytes), material_key, page "
                    "FROM figures")
        figrows = cur.fetchall()
        # The team's material tables have no id column (34-column InDeS
        # schema); order on columns that exist there, like the live DB.
        cur.execute('SELECT property_name, status, origin, figure_id, value_si, '
                    'unit_canonical, flag_reason FROM "Composites_materials" '
                    "ORDER BY extracted_at, origin")
        rows = cur.fetchall()
        cur.execute("SELECT figures_found, figures_mined, figure_rows, vision_calls "
                    "FROM agent_runs ORDER BY id DESC LIMIT 1")
        run_row = cur.fetchone()

    check("cycle harvested exactly 1 figure", m1.get("figures_found") == 1)
    check("cycle mined 1 figure", m1.get("figures_mined") == 1)
    check("cycle inserted 1 figure row", m1.get("figure_rows") == 1)
    check("cycle spent 2 vision calls (1 classify + 1 mine)",
          m1.get("vision_calls") == 2)
    check("text metrics unchanged by figure stage (1 text row inserted)",
          m1.get("rows_inserted") == 1)

    check("figures table has exactly 1 row", len(figrows) == 1)
    if figrows:
        fid, cap, kind, mstat, nbytes, mkey, fpage = figrows[0]
        check("figure caption paired ('Figure 1...')",
              (cap or "").startswith("Figure 1"))
        check("figure classified property_plot + mined",
              kind == "property_plot" and mstat == "mined")
        check("crop bytes stored in Postgres", (nbytes or 0) > 1000)
        check("figure page recorded", fpage == 1)

    fig_rows = [r for r in rows if r[2] == "figure"]
    text_rows = [r for r in rows if r[2] != "figure"]
    check("1 figure row + 1 text row in Composites_materials",
          len(fig_rows) == 1 and len(text_rows) == 1)
    if fig_rows:
        pn, status, origin, fid2, vsi, ucan, flag = fig_rows[0]
        check("figure row is quarantined figure_estimate",
              status == "figure_estimate")
        check("figure row keyed to the figures table",
              bool(fid2) and figrows and fid2 == figrows[0][0])
        check("figure row canonicalized to SI (85 MPa = 8.5e7 Pa)",
              vsi is not None and abs(float(vsi) - 85e6) < 1e3)
    check("run row carries figure counters",
          run_row == (1, 1, 1, 2))

    # --- rerun: done-tracking must spend zero further vision calls ----------
    calls_before = dict(CALLS)
    pdf_path = C.PDF_DIR / "rerun_probe.pdf"
    C.ensure_dirs()
    pdf_path.write_bytes(FAKE_PDF)
    res2 = batch_ingest.process_pdf(
        pdf_path, conn, "test-key", db=pg_mirror,
        figure_opts=batch_ingest.FigureOptions(out_dir=C.FIGURES_DIR))
    check("rerun skipped (sha1 seen)", res2.error == "skipped_seen_sha1")
    check("rerun spent zero vision calls (all figures done)",
          CALLS == calls_before)
    check("rerun inserted no duplicate figure rows", res2.figure_rows == 0)

    # --- publish guard: a figure row arriving status='ok' is downgraded -----
    probe = extraction.PropertyRow(
        material_name="Probe", material_abbreviation="PRB",
        material_class="Composite", material_key="probe",
        section="Mechanical", property_name="Density",
        value="1.3", unit="g/cm3", english="", test_condition="",
        comments="", trade_grade="", manufacturer="", matrix="", fiber="",
        fiber_volume_fraction="", value_raw="1.3", value_num=1.3,
        value_min=None, value_max=None, qualifier="", unit_canonical="g/cm3",
        value_si=1300.0, source_pdf="probe.pdf", source_sha1="probe-sha",
        page=1, source_quote="", status="ok", flag_reason="",
        model="m", prompt_version="p", origin="figure", figure_id="probefig")
    pg_mirror.insert_row(conn, "Composites_materials", probe)
    conn.commit()
    with conn.cursor() as cur:
        cur.execute('SELECT status, flag_reason FROM "Composites_materials" '
                    "WHERE material_key = 'probe'")
        st_, fr_ = cur.fetchone()
    check("publish guard downgraded figure row inserted as ok",
          st_ == "figure_estimate" and "downgraded" in (fr_ or ""))

    # --- origin-scoped promote (the Review Queue SQL) -----------------------
    src = fig_rows and rows and None
    with conn.cursor() as cur:
        # text-only promote must leave the figure row quarantined
        cur.execute('UPDATE "Composites_materials" SET status=\'ok\', '
                    "flag_reason='promoted' WHERE COALESCE(status,'ok') <> 'ok' "
                    "AND COALESCE(origin,'text') = 'text'")
        cur.execute('SELECT count(*) FROM "Composites_materials" '
                    "WHERE origin='figure' AND status <> 'ok'")
        still_q = cur.fetchone()[0]
        # then the include-figures promote blesses them
        cur.execute('UPDATE "Composites_materials" SET status=\'ok\', '
                    "flag_reason='promoted' WHERE COALESCE(status,'ok') <> 'ok'")
        cur.execute('SELECT count(*) FROM "Composites_materials" '
                    "WHERE COALESCE(status,'ok') <> 'ok'")
        left = cur.fetchone()[0]
    conn.commit()
    check("text-only promote left figure rows quarantined", still_q == 2)
    check("include-figures promote clears the queue", left == 0)
finally:
    conn.close()

print("\n%d checks failed" % len(fails))
sys.exit(1 if fails else 0)