test_e2e.py

"""E2E test of the agent cycle against a local Postgres.

Network discovery + Gemini are stubbed; everything else is real code:
orchestrator graph, DOI dedupe, crawler-state persistence, grounding
(verify_against_text on a real generated PDF), unit normalization,
classification, pg_mirror inserts, run/event bookkeeping.
"""
import os
import sys

os.environ.setdefault("AGENT_WORK_DIR", "/tmp/agent_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
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."


def make_pdf() -> bytes:
    doc = fitz.open()
    page = doc.new_page()
    page.insert_text((72, 100), "PEEK CF30 composite datasheet")
    page.insert_text((72, 130), QUOTE)
    data = doc.tobytes()
    doc.close()
    return data


FAKE_PDF = make_pdf()


def fake_search_openalex(query, limit):
    yield pdf_crawler.Candidate(
        title="Tensile behavior of carbon fiber PEEK thermoplastic composite laminates",
        pdf_url="https://example.org/peek_cf30.pdf", source="openalex",
        query=query, doi="10.9999/aim.test.0001", year="2026",
        abstract="thermoplastic composite tensile modulus carbon fiber PEEK datasheet")
    # duplicate DOI from a second query/source — must be deduped pre-download
    yield pdf_crawler.Candidate(
        title="Duplicate DOI variant", pdf_url="https://example.org/peek_dup.pdf",
        source="openalex", query=query, doi="10.9999/aim.test.0001", year="2026",
        abstract="thermoplastic composite tensile carbon fiber")


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),
            extraction.Property(
                section="Mechanical", property_name="Tensile Strength",
                value_raw="99999", value_num=99999.0, unit="MPa",
                test_condition="", source_quote="not in the pdf text at all",
                page=1),
        ])
    return extraction.Extraction(materials=[mat], doc_status="ok")


# --- 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

# --- run two cycles ---------------------------------------------------------
orchestrator.bootstrap()
m1 = orchestrator.run_cycle(trigger="test")
print("cycle 1 metrics:", m1)
m2 = orchestrator.run_cycle(trigger="test")
print("cycle 2 metrics:", m2)

# --- assertions -------------------------------------------------------------
conn = agdb.connect()
fails = []


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


try:
    with conn.cursor() as cur:
        cur.execute('SELECT material_name, status, value_si, unit_canonical, '
                    'source_sha1, prompt_version FROM "Composites_materials"')
        rows = cur.fetchall()
        cur.execute("SELECT count(*) FROM agent_runs WHERE status = 'ok'")
        n_runs = cur.fetchone()[0]
        cur.execute("SELECT doi, ingest_status, rows_inserted FROM agent_doi_seen")
        doi_rows = cur.fetchall()
        cur.execute("SELECT count(*) FROM agent_events")
        n_events = cur.fetchone()[0]

    check("cycle1 downloaded exactly 1 (dup DOI deduped in-cycle)",
          m1.get("downloaded") == 1)
    check("cycle1 inserted 2 rows", m1.get("rows_inserted") == 2)
    check("cycle1 flagged 1 row (out-of-range strength)",
          m1.get("rows_flagged") == 1)
    check("cycle2 downloaded 0 (DOI already in registry)",
          m2.get("downloaded") == 0)
    check("cycle2 skipped >=1 by DOI", m2.get("skipped_doi", 0) >= 1)
    check("2 ok runs recorded", n_runs == 2)
    check("agent_doi_seen has exactly 1 DOI",
          len(doi_rows) == 1 and doi_rows[0][0] == "10.9999/aim.test.0001")
    check("DOI marked ingested with 2 rows",
          doi_rows[0][1] == "ingested" and doi_rows[0][2] == 2)
    check("composite routed to Composites_materials, 2 rows", len(rows) == 2)
    ok_row = [r for r in rows if r[1] == "ok"]
    flg = [r for r in rows if r[1] != "ok"]
    check("grounded modulus row is status ok", len(ok_row) == 1)
    check("modulus value_si ≈ 24.5e9 Pa",
          ok_row and abs(ok_row[0][2] - 24.5e9) < 1e6)
    check("modulus unit_canonical is GPa", ok_row and ok_row[0][3] == "GPa")
    check("bogus strength row flagged", len(flg) == 1 and flg[0][1] in
          ("out_of_range", "unverified"))
    check("prompt_version stamped", all(r[5] for r in rows))
    check("events logged", n_events >= 8)

    state = agdb.get_state(conn, "crawler_state", {})
    check("crawler seen-set persisted to DB",
          "https://example.org/peek_cf30.pdf" in state.get("seen_urls", []))
finally:
    conn.close()

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