generated from john/python-template
229 lines
8.5 KiB
Python
229 lines
8.5 KiB
Python
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.errors import ErrorCategory
|
|
from transcription.services.photos import PhotosService
|
|
from transcription.services.sources import source_mime_type
|
|
from transcription.services.store import SourceStorageError
|
|
from transcription.services.store import StoredSourceFile
|
|
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(
|
|
source.file_path == f"documents/{document.id}/{source.filename}"
|
|
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 source.file_path == f"documents/{result.document_id}/{source.filename}"
|
|
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_photo_stores_file_under_shared_photos_directory(default_session_factory, tmp_path):
|
|
settings = Settings(openrouter_api_key="test-key", upload_dir=tmp_path)
|
|
service = PhotosService(session_factory=default_session_factory, settings=settings)
|
|
|
|
created = await service.create_photo(
|
|
person_id=None,
|
|
filename="portrait.png",
|
|
file_bytes=b"portrait-bytes",
|
|
)
|
|
|
|
stored_path = tmp_path / created.path
|
|
assert stored_path.parent == (tmp_path / "photos")
|
|
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)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_document_job_db_failure_maps_to_internal_category(tmp_path, monkeypatch):
|
|
settings = Settings(openrouter_api_key="test-key", upload_dir=tmp_path)
|
|
stored_path = tmp_path / "documents" / "stored-file.pdf"
|
|
stored_path.parent.mkdir(parents=True, exist_ok=True)
|
|
stored_path.write_bytes(b"payload")
|
|
|
|
async def _fake_store_source_file(**_kwargs) -> StoredSourceFile:
|
|
return StoredSourceFile(
|
|
path=stored_path,
|
|
file_hash="f" * 64,
|
|
file_size_bytes=7,
|
|
)
|
|
|
|
async def _boom_create_records(**_kwargs):
|
|
raise RuntimeError("database unavailable")
|
|
|
|
monkeypatch.setattr("transcription.services.store.store_source_file", _fake_store_source_file)
|
|
monkeypatch.setattr("transcription.services.store._create_document_job_records", _boom_create_records)
|
|
|
|
with pytest.raises(SourceStorageError) as exc_info:
|
|
await create_document_job(
|
|
filename="single-page.pdf",
|
|
file_bytes=b"payload",
|
|
settings=settings,
|
|
)
|
|
|
|
assert exc_info.value.category == ErrorCategory.INFRA_PERSISTENT
|
|
assert exc_info.value.retriable is False
|
|
assert not stored_path.exists()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_job_for_document_db_failure_maps_to_internal_category(async_session, tmp_path, monkeypatch):
|
|
document = Document(id=uuid4(), name="db-failure-doc")
|
|
async_session.add(document)
|
|
await async_session.commit()
|
|
|
|
settings = Settings(openrouter_api_key="test-key", upload_dir=tmp_path)
|
|
stored_path = tmp_path / "documents" / str(document.id) / "stored-file.pdf"
|
|
stored_path.parent.mkdir(parents=True, exist_ok=True)
|
|
stored_path.write_bytes(b"payload")
|
|
|
|
async def _fake_store_source_file(**_kwargs) -> StoredSourceFile:
|
|
return StoredSourceFile(
|
|
path=stored_path,
|
|
file_hash="f" * 64,
|
|
file_size_bytes=7,
|
|
)
|
|
|
|
async def _boom_create_records(**_kwargs):
|
|
raise RuntimeError("database unavailable")
|
|
|
|
monkeypatch.setattr("transcription.services.store.store_source_file", _fake_store_source_file)
|
|
monkeypatch.setattr("transcription.services.store._create_job_for_document_records", _boom_create_records)
|
|
|
|
with pytest.raises(SourceStorageError) as exc_info:
|
|
await create_job_for_document(
|
|
document_id=document.id,
|
|
source_files=[("single-page.pdf", b"payload")],
|
|
session=async_session,
|
|
settings=settings,
|
|
)
|
|
|
|
assert exc_info.value.category == ErrorCategory.INFRA_PERSISTENT
|
|
assert exc_info.value.retriable is False
|
|
assert not stored_path.exists()
|