transcription/tests/services/test_store.py

from pathlib import Path
from uuid import uuid4

import pytest
from sqlmodel import col
from sqlmodel import select

from transcription.config import Settings
from transcription.db.models import Document
from transcription.db.models import Job
from transcription.db.models import JobSource
from transcription.db.models import Source
from transcription.services.people import store_person_portrait
from transcription.services.sources import source_mime_type
from transcription.services.store import SourceStorageError
from transcription.services.store import create_document_job
from transcription.services.store import create_job_for_document
from transcription.services.store import store_source_file


@pytest.mark.asyncio
async def test_create_job_for_document_requires_at_least_one_source(async_session, tmp_path):
    document = Document(id=uuid4(), name="needs-upload")
    async_session.add(document)
    await async_session.commit()

    settings = Settings(openrouter_api_key="test-key", upload_dir=tmp_path)

    with pytest.raises(SourceStorageError):
        await create_job_for_document(
            document_id=document.id,
            source_files=[],
            session=async_session,
            settings=settings,
        )


@pytest.mark.asyncio
async def test_create_job_for_document_sorts_sources_and_creates_links(async_session, tmp_path):
    document = Document(id=uuid4(), name="ordered-upload-doc")
    async_session.add(document)
    await async_session.commit()

    settings = Settings(openrouter_api_key="test-key", upload_dir=tmp_path)

    result = await create_job_for_document(
        document_id=document.id,
        source_files=[
            ("folder/b_page.pdf", b"b"),
            ("folder/A_page.pdf", b"a"),
        ],
        provider="openrouter",
        model="test-model",
        session=async_session,
        settings=settings,
    )

    created_job = await async_session.get(Job, result.job_id)
    assert created_job is not None
    assert created_job.provider == "openrouter"
    assert created_job.model == "test-model"
    assert created_job.prompt_name == "transcribe_document.md"
    assert created_job.user_prompt is not None

    sources = (
        await async_session.exec(
            select(Source).where(Source.document_id == document.id).order_by(col(Source.page_number))
        )
    ).all()
    assert [source.upload_name for source in sources] == ["A_page.pdf", "b_page.pdf"]
    assert all(source.filename.endswith(".pdf") for source in sources)
    assert all("A_page" not in source.filename and "b_page" not in source.filename for source in sources)
    assert all(Path(source.filename).stem == str(source.id) for source in sources)
    assert all(Path(source.file_path).parent == (tmp_path / "documents" / str(document.id)) for source in sources)
    assert [source.file_hash for source in sources] == [
        "ca978112ca1bbdcafac231b39a23dc4da786eff8147c4e72b9807785afee48bb",
        "3e23e8160039594a33894f6564e1b1348bbd7a0088d42c4acb73eeaed59c009d",
    ]
    assert [source.file_size_bytes for source in sources] == [1, 1]

    job_sources = (await async_session.exec(select(JobSource).where(JobSource.job_id == result.job_id))).all()
    assert len(job_sources) == 2
    assert set(result.source_ids) == {job_source.source_id for job_source in job_sources}


@pytest.mark.asyncio
async def test_create_document_job_stores_source_under_document_id_directory(async_session, tmp_path):
    settings = Settings(openrouter_api_key="test-key", upload_dir=tmp_path)

    result = await create_document_job(
        filename="single-page.jpg",
        file_bytes=b"image-bytes",
        session=async_session,
        settings=settings,
    )

    expected_parent = tmp_path / "documents" / str(result.document_id)
    assert result.stored_path.parent == expected_parent
    assert result.stored_path.exists()

    source = (
        await async_session.exec(
            select(Source).where(Source.document_id == result.document_id).order_by(col(Source.page_number))
        )
    ).first()
    assert source is not None
    assert Path(source.filename).stem == str(source.id)
    assert result.stored_path.name == source.filename
    assert Path(source.file_path).parent == expected_parent
    assert source.file_hash == "2c8648d103e3dd7ad87660da0f126a1443b6d21ac1bd3ec000c5e24e2373a90c"
    assert source.file_size_bytes == len(b"image-bytes")

    created_job = await async_session.get(Job, result.job_id)
    assert created_job is not None
    assert created_job.prompt_name == "transcribe_document.md"


@pytest.mark.asyncio
async def test_store_person_portrait_stores_file_under_person_id_directory(tmp_path):
    settings = Settings(openrouter_api_key="test-key", upload_dir=tmp_path)
    person_id = uuid4()

    stored_path = await store_person_portrait(
        person_id=person_id,
        filename="portrait.png",
        file_bytes=b"portrait-bytes",
        settings=settings,
    )

    assert stored_path.parent == (tmp_path / "persons" / str(person_id))
    assert stored_path.exists()


@pytest.mark.parametrize(
    ("filename", "expected_mime_type"),
    [
        ("page.jpg", "image/jpeg"),
        ("page.JPEG", "image/jpeg"),
        ("page.png", "image/png"),
        ("page.tif", "image/tiff"),
        ("page.TIFF", "image/tiff"),
        ("page.pdf", "application/pdf"),
    ],
)
def test_source_mime_type_uses_canonical_source_policy(filename, expected_mime_type):
    assert source_mime_type(filename) == expected_mime_type


@pytest.mark.asyncio
async def test_source_storage_rejects_unsupported_format(tmp_path):
    settings = Settings(openrouter_api_key="test-key", upload_dir=tmp_path)

    with pytest.raises(SourceStorageError):
        await store_source_file(filename="page.txt", file_bytes=b"text", settings=settings)