from pathlib import Path from uuid import uuid4 import pytest 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.store import UploadError from transcription.services.store import create_upload_job from transcription.services.store import create_job_for_document from transcription.services.store import store_person_portrait @pytest.mark.asyncio async def test_create_job_for_document_requires_at_least_one_upload(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(UploadError): await create_job_for_document( document_id=document.id, uploads=[], session=async_session, settings=settings, ) @pytest.mark.asyncio async def test_create_job_for_document_sorts_uploads_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, uploads=[ ("folder/b_page.pdf", b"b"), ("folder/A_page.pdf", b"a"), ], provider="openrouter", model="test-model", prompt_name="transcribe_document.md", 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" sources = ( await async_session.exec( select(Source) .where(Source.document_id == document.id) .order_by(Source.page_number) # pyright: ignore[reportArgumentType] ) ).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) 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_upload_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_upload_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(Source.page_number) # pyright: ignore[reportArgumentType] ) ).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 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 = 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()