generated from john/python-template
121 lines
4.3 KiB
Python
121 lines
4.3 KiB
Python
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()
|