FabledCurator/tests/test_importer_sha_race.py

"""Two workers importing the same picture at once (services/importer.py).

The importer looks a file's sha256 up before inserting its row. When two copies
of one picture are imported side by side, both lookups can find nothing; the
unique index on sha256 then refuses the second insert. The loser files its copy
as a duplicate of the winner, links its post, and leaves the session usable,
instead of failing the import with a dead transaction.

The race is staged the way tests/test_importer_provenance_race.py stages its
own: the loser's sha256 lookup is forced to miss once, as it would if it ran
just before the winner's commit. The loser's near-duplicate cache is empty for
the same reason (it was loaded before the winner's row existed).
"""

import io
import zipfile
from pathlib import Path

import pytest
from PIL import Image
from sqlalchemy import func, select

from backend.app.models import Artist, ImageRecord, ImportSettings, PostAttachment
from backend.app.services.importer import Importer, SkipReason
from backend.app.services.thumbnailer import Thumbnailer

pytestmark = pytest.mark.integration


@pytest.fixture
def layout(tmp_path):
    import_root = tmp_path / "import"
    images_root = tmp_path / "images"
    import_root.mkdir()
    images_root.mkdir()
    return import_root, images_root


def _worker(db_sync, layout) -> Importer:
    import_root, images_root = layout
    settings = db_sync.execute(
        select(ImportSettings).where(ImportSettings.id == 1)
    ).scalar_one()
    return Importer(
        session=db_sync, images_root=images_root, import_root=import_root,
        thumbnailer=Thumbnailer(images_root=images_root), settings=settings,
    )


def _jpeg(path: Path) -> bytes:
    path.parent.mkdir(parents=True, exist_ok=True)
    Image.new("RGB", (320, 240), (40, 90, 160)).save(path, "JPEG")
    return path.read_bytes()


def _copy(src: Path, dst: Path) -> Path:
    dst.parent.mkdir(parents=True, exist_ok=True)
    dst.write_bytes(src.read_bytes())
    return dst


def _lose_the_race(importer: Importer, monkeypatch) -> None:
    """Make `importer` miss its first sha256 lookup and see no near-duplicates."""
    real_execute = importer.session.execute
    missed = [False]

    class _Miss:
        def scalar_one_or_none(self):
            return None

    def execute(stmt, *args, **kwargs):
        sql = str(stmt)
        if not missed[0] and "FROM image_record" in sql and "image_record.sha256 =" in sql:
            missed[0] = True
            return _Miss()
        return real_execute(stmt, *args, **kwargs)

    monkeypatch.setattr(importer.session, "execute", execute)
    monkeypatch.setattr(importer, "_phash_candidates_cache", lambda: [])


def _records(db_sync) -> int:
    return db_sync.execute(select(func.count()).select_from(ImageRecord)).scalar_one()


def test_a_copy_imported_alongside_the_original_is_filed_as_its_duplicate(
    db_sync, layout, monkeypatch,
):
    import_root, images_root = layout
    first = import_root / "Alice" / "set1" / "pic.jpg"
    _jpeg(first)
    second = _copy(first, import_root / "Alice" / "set2" / "pic.jpg")

    winner = _worker(db_sync, layout).import_one(first)
    assert winner.status == "imported"
    kept = Path(db_sync.get(ImageRecord, winner.image_id).path)

    loser = _worker(db_sync, layout)
    _lose_the_race(loser, monkeypatch)
    result = loser.import_one(second)

    assert result.status == "skipped"
    assert result.skip_reason == SkipReason.duplicate_hash
    assert result.image_id == winner.image_id
    monkeypatch.undo()
    assert _records(db_sync) == 1
    assert kept.is_file()
    # The loser's copy in the library was a stray and is gone; only the
    # winner's file remains.
    library_files = [p for p in images_root.rglob("*.jpg") if p.is_file()]
    assert library_files == [kept]


def test_the_same_file_imported_twice_at_once_keeps_its_library_copy(
    db_sync, layout, monkeypatch,
):
    """Both runs write the same library path, so the loser's copy IS the
    winner's file and must not be deleted."""
    import_root, _ = layout
    src = import_root / "Alice" / "pic.jpg"
    _jpeg(src)

    winner = _worker(db_sync, layout).import_one(src)
    kept = Path(db_sync.get(ImageRecord, winner.image_id).path)

    loser = _worker(db_sync, layout)
    _lose_the_race(loser, monkeypatch)
    result = loser.import_one(src)

    assert result.skip_reason == SkipReason.duplicate_hash
    assert kept.is_file()


def test_a_copy_attached_in_place_alongside_the_original_is_its_duplicate(
    db_sync, layout, monkeypatch,
):
    _, images_root = layout
    artist = Artist(name="Bea", slug="bea")
    db_sync.add(artist)
    db_sync.commit()
    first = images_root / "bea" / "patreon" / "post1" / "pic.jpg"
    _jpeg(first)
    second = _copy(first, images_root / "bea" / "patreon" / "post2" / "pic.jpg")

    winner = _worker(db_sync, layout).attach_in_place(first, artist=artist)
    assert winner.status == "imported"

    loser = _worker(db_sync, layout)
    _lose_the_race(loser, monkeypatch)
    result = loser.attach_in_place(second, artist=artist)

    assert result.status == "skipped"
    assert result.skip_reason == SkipReason.duplicate_hash
    assert result.image_id == winner.image_id
    monkeypatch.undo()
    assert _records(db_sync) == 1
    # Removing the loser's file is the caller's step (discard_after_attach).
    assert first.is_file() and second.is_file()


def test_an_archive_whose_member_breaks_the_transaction_is_still_stored(
    db_sync, layout, monkeypatch,
):
    """A member whose import fails mid-flush leaves the transaction dead. The
    member loop rolls it back, so the archive itself is still stored and the
    import ends with a result instead of an exception."""
    import_root, _ = layout
    seed = _worker(db_sync, layout).import_one(_jpeg_at(import_root / "Cy" / "seed.jpg"))
    seed_path = db_sync.get(ImageRecord, seed.image_id).path

    arc = import_root / "Cy" / "set.zip"
    with zipfile.ZipFile(arc, "w") as zf:
        zf.writestr("a.jpg", _jpeg_bytes())

    importer = _worker(db_sync, layout)

    def breaks_the_transaction(*args, **kwargs):
        importer.session.add(ImageRecord(
            path=seed_path, sha256="f" * 64, size_bytes=1, mime="image/jpeg",
            origin="imported_filesystem",
        ))
        importer.session.flush()  # refused: the path is taken

    monkeypatch.setattr(importer, "_import_media", breaks_the_transaction)
    result = importer.import_one(arc)

    assert result.status == "attached"
    stored = db_sync.execute(
        select(func.count()).select_from(PostAttachment)
    ).scalar_one()
    assert stored == 1


def _jpeg_at(path: Path) -> Path:
    _jpeg(path)
    return path


def _jpeg_bytes() -> bytes:
    buf = io.BytesIO()
    Image.new("RGB", (64, 48), (200, 30, 30)).save(buf, "JPEG")
    return buf.getvalue()