stack/tests/llm/test_pdfs.py

"""llm.pdfs — attachments on file for the viewer: comment downloads, bib
attachments, Zotero storage PDFs; name-only resolution inside the roots."""

from __future__ import annotations

import sqlite3
from pathlib import Path

from bib.item import Source
from bib.store import Store
from llm.pdfs import PdfFile, pdf_files, resolve


def _zotero(
    tmp_path: Path, parent_key: str, att_key: str, name: str
) -> tuple[Path, Path]:
    db = tmp_path / "zotero.sqlite"
    con = sqlite3.connect(db)
    con.executescript(
        "CREATE TABLE items (itemID INTEGER PRIMARY KEY, key TEXT);"
        "CREATE TABLE itemAttachments (itemID INTEGER, parentItemID INTEGER, path TEXT);"
        "CREATE TABLE deletedItems (itemID INTEGER);"
    )
    con.execute("INSERT INTO items VALUES (1, ?)", (parent_key,))
    con.execute("INSERT INTO items VALUES (2, ?)", (att_key,))
    con.execute("INSERT INTO items VALUES (3, 'GONE1234')")
    con.execute("INSERT INTO itemAttachments VALUES (2, 1, ?)", (f"storage:{name}",))
    con.execute("INSERT INTO itemAttachments VALUES (3, 1, 'storage:deleted.pdf')")
    con.execute("INSERT INTO deletedItems VALUES (3)")
    con.commit()
    con.close()
    storage = tmp_path / "zstorage"
    (storage / att_key).mkdir(parents=True)
    (storage / att_key / name).write_bytes(b"%PDF-1.4 z")
    (storage / "GONE1234").mkdir()
    (storage / "GONE1234" / "deleted.pdf").write_bytes(b"%PDF-1.4 d")
    return db, storage


def test_all_three_sources_in_order_and_deduped(tmp_path):
    store = Store(":memory:", storage_dir=tmp_path / "bibstorage")
    key = store.upsert(Source(title="Paper", url="https://doi.org/10.1/x"))
    pdf = tmp_path / "paper.pdf"
    pdf.write_bytes(b"%PDF-1.4 b")
    store.attach_file(key, pdf)
    (tmp_path / "notes.docx").write_bytes(b"PK docx")
    store.attach_file(key, tmp_path / "notes.docx")
    zdb, zstore = _zotero(
        tmp_path, key, "ATT00001", "paper.pdf"
    )  # same name as the bib copy → deduped
    files = pdf_files(
        store,
        key,
        zotero_sqlite=zdb,
        zotero_storage=zstore,
        comments_root=tmp_path / "comments",
    )
    assert [(f.name, f.source, f.renderer) for f in files] == [
        ("paper.pdf", "bib", "pdf"),
        ("notes.docx", "bib", "docx"),
    ]
    assert all(f.renderable for f in files)
    assert files[1].media_type.startswith("application/vnd.openxmlformats")
    assert all(isinstance(f, PdfFile) and f.size > 0 for f in files)
    store.close()


def test_zotero_only_item_and_deleted_attachments_skipped(tmp_path):
    store = Store(":memory:", storage_dir=tmp_path / "bibstorage")
    key = store.upsert(Source(title="Zot only", url="https://example.org/z"))
    zdb, zstore = _zotero(tmp_path, key, "ATT00002", "Cassidy - 2025 - RFI.pdf")
    files = pdf_files(store, key, zotero_sqlite=zdb, zotero_storage=zstore)
    assert [(f.name, f.source) for f in files] == [
        ("Cassidy - 2025 - RFI.pdf", "zotero")
    ]
    assert pdf_files(store, "NOPE1234", zotero_sqlite=zdb, zotero_storage=zstore) == []
    assert (
        pdf_files(
            store, key, zotero_sqlite=tmp_path / "missing.sqlite", zotero_storage=zstore
        )
        == []
    )
    store.close()


def test_comment_attachments_from_the_state_tree(tmp_path):
    store = Store(":memory:", storage_dir=tmp_path / "bibstorage")
    key = store.upsert(
        Source(title="c", url="https://www.regulations.gov/comment/CMS-2026-2377-40314")
    )
    d = tmp_path / "comments" / "CMS-2026-2377" / "CMS-2026-2377-40314"
    d.mkdir(parents=True)
    (d / "attachment_1.pdf").write_bytes(b"%PDF-1.4 a1")
    (d / "attachment_2.docx").write_bytes(b"PK")
    (d / "combined.md").write_text("# not an attachment")
    (d / "attachment_1.png").write_bytes(b"png")
    files = pdf_files(store, key, comments_root=tmp_path / "comments")
    (d / "attachment_3.doc").write_bytes(b"\xd0\xcf")
    (d / "attachment_4.txt").write_text("plain")
    files = pdf_files(store, key, comments_root=tmp_path / "comments")
    assert [(f.name, f.source, f.renderer) for f in files] == [
        ("attachment_1.pdf", "comment", "pdf"),
        ("attachment_2.docx", "comment", "docx"),
        ("attachment_3.doc", "comment", None),
        ("attachment_4.txt", "comment", "text"),
    ]
    assert [f.renderable for f in files] == [True, True, False, True]
    # a non-comment url, or a comment with no downloaded dir, yields nothing
    other = store.upsert(Source(title="p", url="https://doi.org/10.1/y"))
    assert pdf_files(store, other, comments_root=tmp_path / "comments") == []
    gone = store.upsert(
        Source(
            title="c2", url="https://www.regulations.gov/comment/CMS-2026-2377-99999"
        )
    )
    assert pdf_files(store, gone, comments_root=tmp_path / "comments") == []
    store.close()


def test_resolve_by_name_inside_roots_only(tmp_path):
    inside = tmp_path / "root" / "K" / "a.pdf"
    inside.parent.mkdir(parents=True)
    inside.write_bytes(b"%PDF")
    outside = tmp_path / "elsewhere" / "b.pdf"
    outside.parent.mkdir()
    outside.write_bytes(b"%PDF")
    files = [PdfFile("a.pdf", inside, 4, "bib"), PdfFile("b.pdf", outside, 4, "bib")]
    roots = (tmp_path / "root",)
    assert resolve(files, "", roots=roots).name == "a.pdf"
    assert resolve(files, "a.pdf", roots=roots).name == "a.pdf"
    assert (
        resolve(files, "b.pdf", roots=roots) is None
    )  # listed, but outside every root
    assert resolve(files, "../a.pdf", roots=roots) is None
    assert resolve([], "", roots=roots) is None