transcription/tests/test_models.py

"""Tests for the V2 SQLModel persistence layer and relationships."""

from datetime import UTC
from datetime import datetime
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