generated from john/python-template
412 lines
16 KiB
Python
412 lines
16 KiB
Python
"""Integration tests for end-to-end upload and worker pipeline behavior."""
|
|
|
|
import io
|
|
from pathlib import Path
|
|
from uuid import uuid4
|
|
|
|
import pytest
|
|
from PIL import Image
|
|
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 JobSourceStatus
|
|
from transcription.db.models import JobStatus
|
|
from transcription.providers.base import ProviderUsage
|
|
from transcription.providers.base import TranscriptionMetadata
|
|
from transcription.providers.base import TranscriptionResult
|
|
from transcription.services import ServiceBundle
|
|
from transcription.services.store import create_document_job
|
|
from transcription.services.store import create_job_for_document
|
|
from transcription.services.workflows import advance_job
|
|
|
|
|
|
async def _attempts_for_job(session, job) -> list[ExecutionAttempt]:
|
|
"""Load execution attempts for a job; V4.7 moved evidence off JobSource."""
|
|
job_source_ids = [job_source.id for job_source in job.job_sources]
|
|
result = await session.exec(
|
|
select(ExecutionAttempt).where(col(ExecutionAttempt.job_source_id).in_(job_source_ids))
|
|
)
|
|
return list(result.all())
|
|
|
|
|
|
def _jpeg_bytes(color: str = "white") -> bytes:
|
|
output = io.BytesIO()
|
|
Image.new("RGB", (2, 2), color=color).save(output, format="JPEG")
|
|
return output.getvalue()
|
|
|
|
|
|
def _build_services(default_session_factory) -> ServiceBundle:
|
|
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),
|
|
)
|
|
return services
|
|
|
|
|
|
@pytest.mark.integration
|
|
class TestPipelineSuccessFlow:
|
|
"""Verify end-to-end success lifecycle behavior."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_upload_then_worker_persists_transcribed_terminal_state(
|
|
self, async_session, default_session_factory, tmp_path: Path, monkeypatch
|
|
):
|
|
"""Upload followed by worker processing persists job transcription and transcribed status."""
|
|
settings = Settings(
|
|
openrouter_api_key="test-key",
|
|
upload_dir=tmp_path,
|
|
default_prompt_name="transcribe_document.md",
|
|
transcription_temperature=0.2,
|
|
transcription_top_p=0.85,
|
|
)
|
|
upload_result = await create_document_job(
|
|
filename="pipeline.jpg",
|
|
file_bytes=_jpeg_bytes(),
|
|
session=async_session,
|
|
settings=settings,
|
|
)
|
|
|
|
async def _fake_transcribe(*, prompt_text: str, image_bytes: bytes, mime_type: str) -> TranscriptionResult:
|
|
_ = (prompt_text, image_bytes, mime_type)
|
|
return TranscriptionResult(
|
|
text="Pipeline transcript",
|
|
provider="openrouter",
|
|
model="test-model",
|
|
prompt_name="transcribe_document.md",
|
|
)
|
|
|
|
async def _fake_transcribe_document_image(
|
|
image_path,
|
|
*,
|
|
prompt_name="transcribe_document.md",
|
|
prompt_text=None,
|
|
temperature=None,
|
|
top_p=None,
|
|
settings=None,
|
|
provider=None,
|
|
) -> TranscriptionResult:
|
|
_ = (image_path, prompt_name, prompt_text, temperature, top_p, settings, provider)
|
|
return TranscriptionResult(
|
|
text="Pipeline transcript",
|
|
provider="openrouter",
|
|
model="test-model",
|
|
prompt_name="transcribe_document.md",
|
|
metadata=TranscriptionMetadata(
|
|
finish_reason="stop",
|
|
usage=ProviderUsage(total_tokens=42),
|
|
),
|
|
raw_api_response={"id": "resp_123", "choices": [{"message": {"content": "Pipeline transcript"}}]},
|
|
)
|
|
|
|
monkeypatch.setattr(
|
|
"transcription.services.workflows.transcribe_document_image",
|
|
_fake_transcribe_document_image,
|
|
)
|
|
|
|
services = _build_services(default_session_factory)
|
|
queued_job = await services.jobs.read_job(job_id=upload_result.job_id, session=async_session)
|
|
processed = queued_job is not None
|
|
if queued_job is not None:
|
|
await advance_job(job=queued_job, services=services, settings=settings, session=async_session)
|
|
job = await services.jobs.read_job(job_id=upload_result.job_id, session=async_session)
|
|
|
|
assert processed is True
|
|
assert job is not None
|
|
assert job.status == JobStatus.TRANSCRIBED
|
|
attempts = await _attempts_for_job(async_session, job)
|
|
assert any(attempt.raw_transcription == "Pipeline transcript" for attempt in attempts)
|
|
assert job.prompt_name == "transcribe_document.md"
|
|
assert job.user_prompt is not None
|
|
assert job.temperature == 0.2
|
|
assert job.top_p == 0.85
|
|
assert any(
|
|
attempt.normalized_metadata is not None
|
|
and attempt.normalized_metadata["finish_reason"] == "stop"
|
|
and attempt.normalized_metadata["usage"] == {"total_tokens": 42}
|
|
for attempt in attempts
|
|
)
|
|
assert any(
|
|
attempt.sdk_response_snapshot
|
|
== {"id": "resp_123", "choices": [{"message": {"content": "Pipeline transcript"}}]}
|
|
for attempt in attempts
|
|
)
|
|
assert all(attempt.error_detail is None for attempt in attempts)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_worker_transcribes_all_sources_for_multi_page_job(
|
|
self,
|
|
async_session,
|
|
default_session_factory,
|
|
tmp_path: Path,
|
|
monkeypatch,
|
|
):
|
|
"""Worker stores transcription output for every source linked to the queued job."""
|
|
settings = Settings(openrouter_api_key="test-key", upload_dir=tmp_path)
|
|
document = Document(id=uuid4(), name="multi-page-document")
|
|
async_session.add(document)
|
|
await async_session.commit()
|
|
|
|
create_result = await create_job_for_document(
|
|
document_id=document.id,
|
|
source_files=[
|
|
("page-01.jpg", _jpeg_bytes("white")),
|
|
("page-02.jpg", _jpeg_bytes("gray")),
|
|
("page-03.jpg", _jpeg_bytes("black")),
|
|
],
|
|
session=async_session,
|
|
settings=settings,
|
|
)
|
|
|
|
async def _fake_transcribe_document_image(
|
|
image_path,
|
|
*,
|
|
prompt_name="transcribe_document.md",
|
|
prompt_text=None,
|
|
temperature=None,
|
|
top_p=None,
|
|
settings=None,
|
|
provider=None,
|
|
) -> TranscriptionResult:
|
|
page_name = Path(image_path).name
|
|
_ = (prompt_name, prompt_text, temperature, top_p, settings, provider)
|
|
return TranscriptionResult(
|
|
text=f"Transcript for {page_name}",
|
|
provider="openrouter",
|
|
model="test-model",
|
|
prompt_name="transcribe_document.md",
|
|
)
|
|
|
|
monkeypatch.setattr(
|
|
"transcription.services.workflows.transcribe_document_image",
|
|
_fake_transcribe_document_image,
|
|
)
|
|
|
|
services = _build_services(default_session_factory)
|
|
queued_job = await services.jobs.read_job(job_id=create_result.job_id, session=async_session)
|
|
assert queued_job is not None
|
|
|
|
await advance_job(job=queued_job, services=services, settings=settings, session=async_session)
|
|
job = await services.jobs.read_job(job_id=create_result.job_id, session=async_session)
|
|
|
|
assert job.status == JobStatus.TRANSCRIBED
|
|
assert len(job.job_sources) == 3
|
|
assert all(job_source.status == JobSourceStatus.TRANSCRIBED for job_source in job.job_sources)
|
|
attempts = await _attempts_for_job(async_session, job)
|
|
assert all(attempt.raw_transcription for attempt in attempts)
|
|
assert all(
|
|
job_source.source is not None and job_source.source.raw_transcription for job_source in job.job_sources
|
|
)
|
|
assert job.prompt_name == "transcribe_document.md"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_worker_marks_partial_success_when_some_sources_fail(
|
|
self,
|
|
async_session,
|
|
default_session_factory,
|
|
tmp_path: Path,
|
|
monkeypatch,
|
|
):
|
|
"""Mixed page outcomes produce PARTIAL_SUCCESS and preserve per-source status."""
|
|
settings = Settings(openrouter_api_key="test-key", upload_dir=tmp_path)
|
|
document = Document(id=uuid4(), name="partial-page-document")
|
|
async_session.add(document)
|
|
await async_session.commit()
|
|
|
|
create_result = await create_job_for_document(
|
|
document_id=document.id,
|
|
source_files=[
|
|
("page-01.jpg", _jpeg_bytes("white")),
|
|
("page-02.jpg", _jpeg_bytes("gray")),
|
|
],
|
|
session=async_session,
|
|
settings=settings,
|
|
)
|
|
|
|
call_count = 0
|
|
|
|
async def _fake_transcribe_document_image(
|
|
image_path,
|
|
*,
|
|
prompt_name="transcribe_document.md",
|
|
prompt_text=None,
|
|
temperature=None,
|
|
top_p=None,
|
|
settings=None,
|
|
provider=None,
|
|
) -> TranscriptionResult:
|
|
nonlocal call_count
|
|
call_count += 1
|
|
_ = (image_path, prompt_name, prompt_text, temperature, top_p, settings, provider)
|
|
if call_count == 2:
|
|
raise RuntimeError("simulated page failure")
|
|
return TranscriptionResult(
|
|
text="Transcript for first page",
|
|
provider="openrouter",
|
|
model="test-model",
|
|
prompt_name="transcribe_document.md",
|
|
)
|
|
|
|
monkeypatch.setattr(
|
|
"transcription.services.workflows.transcribe_document_image",
|
|
_fake_transcribe_document_image,
|
|
)
|
|
|
|
services = _build_services(default_session_factory)
|
|
queued_job = await services.jobs.read_job(job_id=create_result.job_id, session=async_session)
|
|
assert queued_job is not None
|
|
|
|
await advance_job(job=queued_job, services=services, settings=settings, session=async_session)
|
|
job = await services.jobs.read_job(job_id=create_result.job_id, session=async_session)
|
|
|
|
assert job.status == JobStatus.PARTIAL_SUCCESS
|
|
assert len(job.job_sources) == 2
|
|
statuses = {job_source.status for job_source in job.job_sources}
|
|
assert statuses == {JobSourceStatus.TRANSCRIBED, JobSourceStatus.FAILED}
|
|
attempts = await _attempts_for_job(async_session, job)
|
|
assert any(attempt.error_detail is not None for attempt in attempts)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_worker_skips_already_transcribed_sources_on_resubmit(
|
|
self,
|
|
async_session,
|
|
default_session_factory,
|
|
tmp_path: Path,
|
|
monkeypatch,
|
|
):
|
|
"""Queued jobs only process non-transcribed JobSource records."""
|
|
settings = Settings(openrouter_api_key="test-key", upload_dir=tmp_path)
|
|
document = Document(id=uuid4(), name="resubmit-filter-document")
|
|
async_session.add(document)
|
|
await async_session.commit()
|
|
|
|
create_result = await create_job_for_document(
|
|
document_id=document.id,
|
|
source_files=[
|
|
("page-01.jpg", _jpeg_bytes("white")),
|
|
("page-02.jpg", _jpeg_bytes("gray")),
|
|
],
|
|
session=async_session,
|
|
settings=settings,
|
|
)
|
|
|
|
services = _build_services(default_session_factory)
|
|
job = await services.jobs.read_job(job_id=create_result.job_id, session=async_session)
|
|
page_one = next(js for js in job.job_sources if js.source is not None and js.source.page_number == 1)
|
|
page_two = next(js for js in job.job_sources if js.source is not None and js.source.page_number == 2)
|
|
|
|
page_one.status = JobSourceStatus.TRANSCRIBED
|
|
page_two.status = JobSourceStatus.PENDING
|
|
await services.sources.update_job_source(job_source=page_one, session=async_session)
|
|
await services.sources.update_job_source(job_source=page_two, session=async_session)
|
|
await services.jobs.update_job_state(job_id=job.id, status=JobStatus.QUEUED, session=async_session)
|
|
await async_session.commit()
|
|
|
|
call_count = 0
|
|
|
|
async def _fake_transcribe_document_image(
|
|
image_path,
|
|
*,
|
|
prompt_name="transcribe_document.md",
|
|
prompt_text=None,
|
|
temperature=None,
|
|
top_p=None,
|
|
settings=None,
|
|
provider=None,
|
|
) -> TranscriptionResult:
|
|
nonlocal call_count
|
|
_ = (image_path, prompt_name, prompt_text, temperature, top_p, settings, provider)
|
|
call_count += 1
|
|
return TranscriptionResult(
|
|
text="new transcript",
|
|
provider="openrouter",
|
|
model="test-model",
|
|
prompt_name="transcribe_document.md",
|
|
)
|
|
|
|
monkeypatch.setattr(
|
|
"transcription.services.workflows.transcribe_document_image",
|
|
_fake_transcribe_document_image,
|
|
)
|
|
|
|
queued_job = await services.jobs.read_job(job_id=create_result.job_id, session=async_session)
|
|
assert queued_job is not None
|
|
await advance_job(job=queued_job, services=services, settings=settings, session=async_session)
|
|
|
|
refreshed = await services.jobs.read_job(job_id=create_result.job_id, session=async_session)
|
|
assert call_count == 1
|
|
statuses = {js.status for js in refreshed.job_sources}
|
|
assert statuses == {JobSourceStatus.TRANSCRIBED}
|
|
|
|
|
|
@pytest.mark.integration
|
|
class TestPipelineFailureFlow:
|
|
"""Verify end-to-end failure lifecycle behavior."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_upload_then_worker_persists_failed_terminal_state(
|
|
self,
|
|
async_session,
|
|
default_session_factory,
|
|
tmp_path: Path,
|
|
monkeypatch,
|
|
):
|
|
"""Upload followed by worker processing persists error detail and failed status on the job."""
|
|
settings = Settings(openrouter_api_key="test-key", upload_dir=tmp_path)
|
|
upload_result = await create_document_job(
|
|
filename="pipeline.jpg",
|
|
file_bytes=_jpeg_bytes(),
|
|
session=async_session,
|
|
settings=settings,
|
|
)
|
|
|
|
async def _fake_transcribe_document_image(
|
|
image_path,
|
|
*,
|
|
prompt_name="transcribe_document.md",
|
|
prompt_text=None,
|
|
temperature=None,
|
|
top_p=None,
|
|
settings=None,
|
|
provider=None,
|
|
) -> TranscriptionResult:
|
|
_ = (image_path, prompt_name, prompt_text, temperature, top_p, settings, provider)
|
|
raise RuntimeError("pipeline provider failure")
|
|
|
|
monkeypatch.setattr(
|
|
"transcription.services.workflows.transcribe_document_image",
|
|
_fake_transcribe_document_image,
|
|
)
|
|
|
|
services = _build_services(default_session_factory)
|
|
queued_job = await services.jobs.read_job(job_id=upload_result.job_id, session=async_session)
|
|
processed = queued_job is not None
|
|
if queued_job is not None:
|
|
await advance_job(job=queued_job, services=services, settings=settings, session=async_session)
|
|
job = await services.jobs.read_job(job_id=upload_result.job_id, session=async_session)
|
|
|
|
assert processed is True
|
|
assert job is not None
|
|
assert job.status == JobStatus.FAILED
|
|
attempts = await _attempts_for_job(async_session, job)
|
|
assert all(attempt.raw_transcription is None for attempt in attempts)
|
|
assert any(attempt.error_detail is not None for attempt in attempts)
|
|
error_detail = next(attempt.error_detail for attempt in attempts if attempt.error_detail is not None)
|
|
assert "pipeline provider failure" in error_detail
|
|
assert "[internal_unexpected_error]" in error_detail
|
|
assert "error_id=" in error_detail
|