stack/tests/zot/fetch/test_attach.py

from __future__ import annotations

from pathlib import Path

import pytest

from tests.zot.fetch.conftest import mk_item
from zot.db import Db
from zot.fetch.attach import (
    CONTENT_TYPES,
    append_extra,
    attach_file,
    has_full_text,
    sudo_copier,
)


def _py_copier(src: Path, dest_dir: Path) -> None:
    dest_dir.mkdir(parents=True, exist_ok=True)
    (dest_dir / src.name).write_bytes(src.read_bytes())


def _att(db: Db, parent: int, ctype: str) -> None:
    db.add_attachment(parent, content_type=ctype, path="storage:x")


@pytest.mark.parametrize(
    "ctypes, expected",
    [
        ([], False),
        (["text/html"], False),
        (["application/pdf"], True),
        (["application/epub+zip"], True),
        (["text/html", "application/pdf"], True),
    ],
)
def test_has_full_text_truth_table(zdb: Db, ctypes, expected):
    parent = mk_item(zdb, title="P")
    for c in ctypes:
        _att(zdb, parent, c)
    assert has_full_text(zdb, parent) is expected


def test_has_full_text_ignores_other_parents(zdb: Db):
    a = mk_item(zdb, title="A")
    b = mk_item(zdb, title="B")
    _att(zdb, a, "application/pdf")
    assert has_full_text(zdb, b) is False


def test_attach_file_writes_row_tags_and_extra(zdb: Db, tmp_path: Path):
    parent = mk_item(zdb, title="Paper", extra="PMID: 1")
    src = tmp_path / "Paper.pdf"
    src.write_bytes(b"%PDF-1.4 x")
    storage = tmp_path / "storage"
    att_id = attach_file(
        zdb,
        parent,
        src,
        source="libgen",
        md5="a" * 32,
        storage_dir=storage,
        copier=_py_copier,
    )
    row = zdb.con.execute(
        "SELECT parentItemID, linkMode, contentType, path FROM itemAttachments "
        "WHERE itemID = ?",
        (att_id,),
    ).fetchone()
    assert tuple(row) == (parent, 0, "application/pdf", "storage:Paper.pdf")
    assert zdb.get_field(att_id, "title") == "Paper.pdf"
    att_key = zdb.get_item(att_id)["key"]
    assert (storage / att_key / "Paper.pdf").read_bytes() == b"%PDF-1.4 x"
    assert "source:libgen" in zdb.get_tags(parent)
    assert zdb.get_field(parent, "extra") == "PMID: 1\nmd5: " + "a" * 32
    assert has_full_text(zdb, parent) is True


def test_attach_file_epub_content_type(zdb: Db, tmp_path: Path):
    parent = mk_item(zdb, type_name="book", title="Book")
    src = tmp_path / "Book.epub"
    src.write_bytes(b"PK\x03\x04")
    att_id = attach_file(
        zdb, parent, src, source="annas", storage_dir=tmp_path / "s", copier=_py_copier
    )
    ctype = zdb.con.execute(
        "SELECT contentType FROM itemAttachments WHERE itemID = ?", (att_id,)
    ).fetchone()[0]
    assert ctype == CONTENT_TYPES[".epub"]
    assert zdb.get_field(parent, "extra") is None  # no md5 → extra untouched


def test_attach_file_rejects_unknown_extension(zdb: Db, tmp_path: Path):
    parent = mk_item(zdb, title="P")
    src = tmp_path / "x.mobi"
    src.write_bytes(b"x")
    with pytest.raises(ValueError, match="mobi"):
        attach_file(
            zdb, parent, src, source="libgen", storage_dir=tmp_path, copier=_py_copier
        )


def test_append_extra_is_idempotent_and_keeps_lines(zdb: Db):
    item = mk_item(zdb, title="P", extra="PMID: 1")
    append_extra(zdb, item, "md5: abc")
    append_extra(zdb, item, "md5: abc")
    assert zdb.get_field(item, "extra") == "PMID: 1\nmd5: abc"


def test_append_extra_on_empty_extra(zdb: Db):
    item = mk_item(zdb, title="P")
    append_extra(zdb, item, "md5: abc")
    assert zdb.get_field(item, "extra") == "md5: abc"


def test_sudo_copier_runs_mkdir_cp_chown(monkeypatch, tmp_path: Path):
    calls: list[list[str]] = []

    def fake_run(cmd, check):
        calls.append(cmd)

    monkeypatch.setattr("zot.fetch.attach.subprocess.run", fake_run)
    sudo_copier(tmp_path / "a.pdf", tmp_path / "storage" / "KEY")
    dest = tmp_path / "storage" / "KEY"
    assert calls == [
        ["sudo", "mkdir", "-p", str(dest)],
        ["sudo", "cp", str(tmp_path / "a.pdf"), str(dest / "a.pdf")],
        ["sudo", "chown", "-R", "100999:100999", str(dest)],
    ]


def test_attach_file_writes_rows_before_copying(zdb: Db, tmp_path: Path):
    """A failed copy must leave rows the caller can roll back, never an
    orphan storage dir with no row."""
    parent = mk_item(zdb, title="P")
    src = tmp_path / "P.pdf"
    src.write_bytes(b"%PDF-1.4 x")
    seen = {}

    def spy_copier(s: Path, d: Path) -> None:
        seen["rows"] = zdb.con.execute(
            "SELECT COUNT(*) FROM itemAttachments WHERE parentItemID = ?", (parent,)
        ).fetchone()[0]
        raise OSError("disk full")

    with pytest.raises(OSError):
        attach_file(
            zdb, parent, src, source="libgen", storage_dir=tmp_path, copier=spy_copier
        )
    assert seen["rows"] == 1
    zdb.con.rollback()
    assert has_full_text(zdb, parent) is False