"""Tests for the V2 SQLModel persistence layer and relationships.""" from uuid import UUID import pytest from sqlalchemy.exc import IntegrityError from transcription.db.models import Document from transcription.db.models import DocumentPerson from transcription.db.models import DocumentPersonRole from transcription.db.models import Job from transcription.db.models import JobSource from transcription.db.models import JobSourceStatus from transcription.db.models import JobStatus from transcription.db.models import Person from transcription.db.models import Source def _make_document(**overrides) -> Document: defaults = { "name": "letter bundle", "document_type": "letter", "notes": "Family correspondence", } defaults.update(overrides) return Document(**defaults) def _persist_document(session) -> Document: document = _make_document() session.add(document) session.commit() session.refresh(document) return document def _persist_person(session, **overrides) -> Person: defaults = {"full_name": "Ada Lovelace"} defaults.update(overrides) person = Person(**defaults) session.add(person) session.commit() session.refresh(person) return person def _persist_job(session, document: Document) -> Job: job = Job(document_id=document.id) session.add(job) session.commit() session.refresh(job) return job def _persist_source(session, document: Document, *, page_number: int = 1, **overrides) -> Source: defaults = { "document_id": document.id, "page_number": page_number, "upload_name": "letter.jpg", "filename": "stored-letter.jpg", "file_path": "/uploads/stored-letter.jpg", "raw_transcription": "Original machine text", } defaults.update(overrides) source = Source(**defaults) session.add(source) session.commit() session.refresh(source) return source def _persist_job_source(session, job: Job, source: Source, **overrides) -> JobSource: defaults = { "job_id": job.id, "source_id": source.id, "status": JobSourceStatus.PENDING, } defaults.update(overrides) job_source = JobSource(**defaults) session.add(job_source) session.commit() session.refresh(job_source) return job_source class TestDocumentModel: def test_can_be_persisted(self, session): document = _persist_document(session) fetched = session.get(Document, document.id) assert fetched is not None assert fetched.name == "letter bundle" def test_defaults_are_populated(self, session): document = _persist_document(session) assert isinstance(document.id, UUID) assert document.created_at is not None assert document.updated_at is not None class TestJobModel: def test_can_be_created_for_document(self, session): document = _persist_document(session) job = _persist_job(session, document) fetched = session.get(Job, job.id) assert fetched is not None assert fetched.document_id == document.id def test_defaults_are_populated(self, session): document = _persist_document(session) job = _persist_job(session, document) assert job.status == JobStatus.QUEUED assert job.retry_count == 0 assert job.date_created is not None assert job.date_updated is not None def test_transitions_to_completed(self, session): document = _persist_document(session) job = _persist_job(session, document) job.status = JobStatus.PROCESSING session.add(job) session.commit() session.refresh(job) job.status = JobStatus.COMPLETED session.add(job) session.commit() session.refresh(job) assert job.status == JobStatus.COMPLETED class TestSourceModel: def test_can_be_created_for_document(self, session): document = _persist_document(session) source = _persist_source(session, document) fetched = session.get(Source, source.id) assert fetched is not None assert fetched.document_id == document.id assert fetched.page_number == 1 assert fetched.date_uploaded is not None def test_revised_text_is_supported(self, session): document = _persist_document(session) source = _persist_source(session, document, revised_text="Edited output") fetched = session.get(Source, source.id) assert fetched is not None assert fetched.revised_text == "Edited output" class TestPersonAndDocumentPersonModel: def test_document_person_role_is_unique_per_document_person(self, session): document = _persist_document(session) person = _persist_person(session) first = DocumentPerson(document_id=document.id, person_id=person.id, role=DocumentPersonRole.AUTHOR) session.add(first) session.commit() duplicate = DocumentPerson(document_id=document.id, person_id=person.id, role=DocumentPersonRole.AUTHOR) session.add(duplicate) with pytest.raises(IntegrityError): session.commit() class TestJobSourceModel: def test_job_source_persists_json_payloads(self, session): document = _persist_document(session) job = _persist_job(session, document) source = _persist_source(session, document) job_source = _persist_job_source( session, job, source, raw_transcription="Page transcript", ai_metadata={"confidence": 0.91, "boxes": [{"x": 1, "y": 2}]}, raw_api_response={"provider": "test"}, ) fetched = session.get(JobSource, job_source.id) assert fetched is not None assert fetched.status == JobSourceStatus.PENDING assert fetched.ai_metadata == {"confidence": 0.91, "boxes": [{"x": 1, "y": 2}]} assert fetched.raw_api_response == {"provider": "test"} class TestRelationships: def test_document_exposes_jobs_sources_and_people(self, session): document = _persist_document(session) _persist_job(session, document) _persist_source(session, document) person = _persist_person(session) link = DocumentPerson(document_id=document.id, person_id=person.id, role=DocumentPersonRole.AUTHOR) session.add(link) session.commit() session.refresh(document) assert len(document.jobs) == 1 assert len(document.sources) == 1 assert len(document.document_people) == 1 def test_document_exposes_author_via_role_filtered_relation(self, session): document = _persist_document(session) author = _persist_person(session, full_name="Author Person") recipient = _persist_person(session, full_name="Recipient Person") session.add(DocumentPerson(document_id=document.id, person_id=author.id, role=DocumentPersonRole.AUTHOR)) session.add(DocumentPerson(document_id=document.id, person_id=recipient.id, role=DocumentPersonRole.RECIPIENT)) session.commit() session.refresh(document) assert [person.full_name for person in document.authors] == ["Author Person"] assert document.author is not None assert document.author.full_name == "Author Person" def test_person_exposes_authored_documents_via_role_filtered_relation(self, session): authored_document = _make_document(name="Authored Doc") recipient_only_document = _make_document(name="Recipient Doc") session.add(authored_document) session.add(recipient_only_document) session.commit() session.refresh(authored_document) session.refresh(recipient_only_document) person = _persist_person(session, full_name="Dual Role Person") session.add( DocumentPerson( document_id=authored_document.id, person_id=person.id, role=DocumentPersonRole.AUTHOR, ) ) session.add( DocumentPerson( document_id=recipient_only_document.id, person_id=person.id, role=DocumentPersonRole.RECIPIENT, ) ) session.commit() session.refresh(person) assert [document.name for document in person.authored_documents] == ["Authored Doc"]