transcription/tests/services/test_workflows_reliability.py

"""Reliability tests for worker workflow timeout behavior."""

import asyncio
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 ExecutionAttempt
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 Source
from transcription.providers.base import TranscriptionResult
from transcription.services import ServiceBundle
from transcription.services.workflows import process_queued_job


@pytest.mark.integration
class TestWorkflowReliability:
    """Verify timeout and terminal-state reliability behavior."""

    @pytest.mark.asyncio
    async def test_process_queued_job_timeout_marks_job_failed(self, default_session_factory, monkeypatch):
        """Provider timeout transitions a queued job to failed with error detail."""
        services = ServiceBundle()
        object.__setattr__(
            services,
            "documents",
            services.documents.__class__(session_factory=default_session_factory),
        )
        object.__setattr__(
            services,
            "jobs",
            services.jobs.__class__(session_factory=default_session_factory),
        )
        object.__setattr__(
            services,
            "sources",
            services.sources.__class__(session_factory=default_session_factory),
        )

        async with services.jobs._session_scope() as session:
            document = Document(id=uuid4(), name="timeout-doc")
            session.add(document)
            await session.flush()

            job = Job(document_id=document.id, status=JobStatus.QUEUED)
            session.add(job)
            await session.flush()

            source = Source(
                document_id=document.id,
                page_number=1,
                upload_name="timeout.jpg",
                filename="timeout.jpg",
                file_path=str(Path("tests/fixtures/images/real/Book Two - page 02.jpg")),
                file_hash="c" * 64,
                file_size_bytes=1,
            )
            session.add(source)
            await session.flush()
            session.add(JobSource(job_id=job.id, source_id=source.id, status=JobSourceStatus.PENDING))
            await session.commit()

            loaded = await services.jobs.read_job(job_id=job.id, session=session)

        async def _never_returns(
            image_path,
            *,
            prompt_name="transcribe_document.md",
            prompt_text=None,
            temperature=None,
            top_p=None,
            settings=None,
            provider=None,
        ):
            _ = (image_path, prompt_name, prompt_text, temperature, top_p, settings, provider)
            raise TimeoutError("simulated provider timeout")

        monkeypatch.setattr("transcription.services.workflows.transcribe_document_image", _never_returns)

        timeout_settings = Settings(openrouter_api_key="test-key", worker_provider_timeout_seconds=20.0)
        result = await process_queued_job(job=loaded, services=services, settings=timeout_settings)

        assert result is not None
        assert result.status == JobStatus.FAILED

        async with services.jobs._session_scope() as session:
            attempts = (
                await session.execute(
                    select(ExecutionAttempt).where(
                        col(ExecutionAttempt.job_source_id).in_([js.id for js in result.job_sources])
                    )
                )
            ).scalars().all()
        error_detail = next(attempt.error_detail for attempt in attempts if attempt.error_detail is not None)
        assert "timed out" in error_detail.lower()
        assert "20.0s" in error_detail

    @pytest.mark.asyncio
    async def test_completed_page_is_committed_before_next_provider_call_finishes(
        self,
        default_session_factory,
        monkeypatch,
    ):
        services = ServiceBundle(
            documents=ServiceBundle().documents.__class__(session_factory=default_session_factory),
            jobs=ServiceBundle().jobs.__class__(session_factory=default_session_factory),
            sources=ServiceBundle().sources.__class__(session_factory=default_session_factory),
            people=ServiceBundle().people.__class__(session_factory=default_session_factory),
        )
        async with services.jobs._session_scope() as session:
            document = Document(id=uuid4(), name="durability-doc")
            session.add(document)
            await session.flush()
            job = Job(document_id=document.id, status=JobStatus.QUEUED)
            session.add(job)
            await session.flush()
            for page_number in (1, 2):
                source = Source(
                    document_id=document.id,
                    page_number=page_number,
                    upload_name=f"page-{page_number}.jpg",
                    filename=f"page-{page_number}.jpg",
                    file_path=str(Path("tests/fixtures/images/real/Book Two - page 02.jpg")),
                    file_hash=str(page_number) * 64,
                    file_size_bytes=1,
                )
                session.add(source)
                await session.flush()
                session.add(JobSource(job_id=job.id, source_id=source.id))
            await session.commit()
            loaded = await services.jobs.read_job(job_id=job.id, session=session)

        second_started = asyncio.Event()
        release_second = asyncio.Event()
        call_count = 0

        async def _transcribe(image_path, **kwargs):
            nonlocal call_count
            _ = (image_path, kwargs)
            call_count += 1
            if call_count == 2:
                second_started.set()
                await release_second.wait()
            return TranscriptionResult(text=f"page {call_count}", provider="fixture", model="model")

        monkeypatch.setattr("transcription.services.workflows.transcribe_document_image", _transcribe)
        task = asyncio.create_task(process_queued_job(job=loaded, services=services))
        await asyncio.wait_for(second_started.wait(), timeout=2)

        attempts = await services.sources.list_execution_attempts(job_id=job.id)
        assert len(attempts) == 1
        assert attempts[0].raw_transcription == "page 1"

        release_second.set()
        result = await task
        assert result is not None
        assert result.status == JobStatus.TRANSCRIBED