photools3/test/test_processing.py

from pathlib import Path
from unittest.mock import MagicMock, patch

import pytest
from PIL import Image

from photools3.models import (
    LlmResponseCaption,
    LlmResponseCaptionsAndKeywords,
    LlmResponseKeywords,
)
from photools3.processing import (
    ProcessingMode,
    _create_session,
    process_batch,
    process_photo,
)


@pytest.fixture
def image_path(tmp_path):
    path = tmp_path / "photo.png"
    Image.new("RGB", (20, 20)).save(path)
    return path


@pytest.mark.parametrize(
    ("mode", "response", "expected_metadata"),
    [
        (
            ProcessingMode.KEYWORDS,
            LlmResponseKeywords(tags=["architektur"]),
            {"caption": None, "keywords": ["architektur"]},
        ),
        (
            ProcessingMode.CAPTIONS,
            LlmResponseCaption(caption="Ein Haus"),
            {"caption": "Ein Haus", "keywords": None},
        ),
        (
            ProcessingMode.BOTH,
            LlmResponseCaptionsAndKeywords(caption="Ein Haus", tags=["architektur"]),
            {"caption": "Ein Haus", "keywords": ["architektur"]},
        ),
    ],
)
def test_process_photo_writes_only_requested_metadata(
    image_path, mode, response, expected_metadata
):
    session = MagicMock()
    session.call.return_value = response

    with (
        patch("photools3.processing._create_session", return_value=session),
        patch("photools3.processing.MetadataWriter") as writer_class,
    ):
        result = process_photo(image_path, mode)

    assert result == response
    session.call.assert_called_once()
    session.close.assert_called_once()
    writer_class.assert_called_once_with(image_path)
    writer_class.return_value.write.assert_called_once_with(**expected_metadata)


def test_process_photo_does_not_close_a_caller_owned_session(image_path):
    session = MagicMock()
    session.call.return_value = LlmResponseCaption(caption="Ein Haus")

    with patch("photools3.processing.MetadataWriter"):
        process_photo(image_path, ProcessingMode.CAPTIONS, session=session)

    session.close.assert_not_called()


@pytest.mark.parametrize(
    ("mode", "constructor"),
    [
        (ProcessingMode.KEYWORDS, "LlmSessionKeywords"),
        (ProcessingMode.CAPTIONS, "LlmSessionCaption"),
        (ProcessingMode.BOTH, "LlmSessionCaptionsAndKeywords"),
    ],
)
def test_create_session_selects_mode_specific_response_model(mode, constructor):
    with patch(f"photools3.processing.{constructor}") as session_class:
        assert _create_session(mode) is session_class.return_value

    session_class.assert_called_once_with()


def test_process_batch_is_sequential_reuses_session_and_continues_after_error():
    active_photos = 0
    max_active_photos = 0

    class TrackingPhoto:
        def __init__(self, path):
            self.path = path

        def __enter__(self):
            nonlocal active_photos, max_active_photos
            active_photos += 1
            max_active_photos = max(max_active_photos, active_photos)
            return self

        def __exit__(self, exc_type, exc_value, traceback):
            nonlocal active_photos
            active_photos -= 1

    session = MagicMock()
    response = LlmResponseCaptionsAndKeywords(caption="Ein Haus", tags=["architektur"])
    session.call.side_effect = [response, RuntimeError("temporary error"), response]
    paths = [Path("one.jpg"), Path("two.jpg"), Path("three.jpg")]

    with (
        patch("photools3.processing._create_session", return_value=session),
        patch("photools3.processing.Photo", TrackingPhoto),
        patch("photools3.processing.MetadataWriter") as writer_class,
    ):
        outcomes = process_batch(paths, ProcessingMode.BOTH)

    assert [outcome.succeeded for outcome in outcomes] == [True, False, True]
    assert [outcome.path for outcome in outcomes] == paths
    assert outcomes[1].error == "RuntimeError: temporary error"
    assert max_active_photos == 1
    assert active_photos == 0
    assert session.call.call_count == 3
    assert all(len(call.args[0]) == 1 for call in session.call.call_args_list)
    assert writer_class.return_value.write.call_count == 2
    session.close.assert_called_once()


def test_process_batch_flushes_after_configured_attempt_count(monkeypatch):
    monkeypatch.setattr("photools3.processing.settings.llm_flush_every_photos", 2)

    class FakePhoto:
        def __init__(self, path):
            self.path = path

        def __enter__(self):
            return self

        def __exit__(self, exc_type, exc_value, traceback):
            return None

    response = LlmResponseCaptionsAndKeywords(caption="Ein Haus", tags=["architektur"])
    session = MagicMock()
    session.call.side_effect = [
        response,
        RuntimeError("temporary error"),
        response,
        response,
        response,
    ]
    paths = [Path(f"{i}.jpg") for i in range(5)]

    with (
        patch("photools3.processing._create_session", return_value=session),
        patch("photools3.processing.Photo", FakePhoto),
        patch("photools3.processing.MetadataWriter"),
    ):
        outcomes = process_batch(paths, ProcessingMode.BOTH)

    assert len(outcomes) == len(paths)
    assert session.flush_model.call_count == 2
    session.close.assert_called_once()