from __future__ import annotations import hashlib import uuid from datetime import UTC, datetime, timedelta from pathlib import Path from unittest.mock import MagicMock import pytest from app.core.config import Settings from app.persistence.document_workflows import ( FAIL_PARSE_SQL, FINALIZE_JOB_SQL, LOAD_SOURCE_SQL, SELECT_BLOCKS_SQL, SELECT_CHUNKS_SQL, SELECT_MANIFEST_ITEMS_SQL, SELECT_PAGES_SQL, SELECT_VERSION_SQL, UPDATE_DOCUMENT_SQL, ArtifactConflictError, DocumentSource, InvalidDocumentJobError, PostgresDocumentWorkflowRepository, _canonical_json, plan_artifact, ) from app.persistence.job_queue import BackgroundJob, JobLease, LeaseLostError from app.services.document_ingestion import CloudTextPolicy, ingest_document NOW = datetime(2026, 7, 13, 8, 0, tzinfo=UTC) JOB_ID = uuid.UUID("10000000-0000-0000-0000-000000000001") UPLOAD_ID = uuid.UUID("20000000-0000-0000-0000-000000000002") DOCUMENT_ID = uuid.UUID("30000000-0000-0000-0000-000000000003") KB_ID = uuid.UUID("40000000-0000-0000-0000-000000000004") SCOPE_ID = uuid.UUID("50000000-0000-0000-0000-000000000005") STORAGE_KEY = uuid.UUID("60000000-0000-0000-0000-000000000006") LEASE_TOKEN = uuid.UUID("70000000-0000-0000-0000-000000000007") RAW_HASH = hashlib.sha256(b"source").hexdigest() def _job(**changes: object) -> BackgroundJob: values: dict[str, object] = { "id": JOB_ID, "job_type": "PARSE_DOCUMENT", "required_capability": "document_parse", "resource_type": "document", "resource_id": DOCUMENT_ID, "idempotency_key": f"parse-document:{DOCUMENT_ID}", "payload": {"upload_id": str(UPLOAD_ID), "document_id": str(DOCUMENT_ID)}, "stage": "PENDING", "progress": 0, "priority": 0, "attempt": 1, "max_attempts": 3, "run_after": NOW, "lease_until": NOW + timedelta(seconds=60), "created_at": NOW, "updated_at": NOW, "lease": JobLease(JOB_ID, "worker-documents", LEASE_TOKEN), } values.update(changes) return BackgroundJob(**values) # type: ignore[arg-type] def _source() -> DocumentSource: return DocumentSource( upload_id=UPLOAD_ID, document_id=DOCUMENT_ID, knowledge_base_id=KB_ID, access_scope_id=SCOPE_ID, filename="source.txt", mime_type="text/plain", storage_key=STORAGE_KEY, byte_size=6, raw_sha256=RAW_HASH, ) def _source_row() -> dict[str, object]: return { "upload_id": UPLOAD_ID, "document_id": DOCUMENT_ID, "knowledge_base_id": KB_ID, "access_scope_id": SCOPE_ID, "original_filename": "source.txt", "declared_mime_type": "text/plain", "storage_key": STORAGE_KEY, "actual_size": 6, "actual_sha256": RAW_HASH, "document_filename": "source.txt", "document_mime_type": "text/plain", "document_raw_sha256": RAW_HASH, "document_storage_key": str(STORAGE_KEY), } def _settings(tmp_path: Path) -> Settings: secret = tmp_path / "postgres-password" secret.write_text("synthetic-password", encoding="utf-8") return Settings(postgres_password_file=secret) def _repository_with_one( tmp_path: Path, row: dict[str, object] | None ) -> tuple[PostgresDocumentWorkflowRepository, MagicMock, MagicMock]: cursor = MagicMock() cursor.fetchone.return_value = row connection = MagicMock() connection.__enter__.return_value = connection connection.execute.return_value = cursor factory = MagicMock(return_value=connection) repository = PostgresDocumentWorkflowRepository(_settings(tmp_path), connection_factory=factory) return repository, connection, factory def test_load_source_checks_full_active_fence_job_payload_and_storage_binding( tmp_path: Path, ) -> None: repository, connection, _ = _repository_with_one(tmp_path, _source_row()) source = repository.load_source(_job()) assert source == _source() statement, parameters = connection.execute.call_args.args assert statement == LOAD_SOURCE_SQL assert "job.status = 'RUNNING'" in statement assert "job.lease_owner = %s" in statement assert "job.lease_token = %s" in statement assert "job.lease_until >= now()" in statement assert "upload.actual_sha256 = document.raw_sha256" in statement assert "upload.storage_key::text = document.storage_key" in statement assert parameters == ( JOB_ID, "worker-documents", LEASE_TOKEN, DOCUMENT_ID, UPLOAD_ID, DOCUMENT_ID, ) def test_invalid_payload_fails_before_database_and_expired_fence_has_no_source( tmp_path: Path, ) -> None: invalid_repository, _, invalid_factory = _repository_with_one(tmp_path, _source_row()) with pytest.raises(InvalidDocumentJobError): invalid_repository.load_source(_job(payload={"document_id": str(DOCUMENT_ID)})) with pytest.raises(InvalidDocumentJobError): invalid_repository.load_source( _job(lease=JobLease(uuid.uuid4(), "worker-documents", LEASE_TOKEN)) ) invalid_factory.assert_not_called() expired_repository, _, _ = _repository_with_one(tmp_path, None) with pytest.raises(LeaseLostError): expired_repository.load_source(_job()) def test_terminal_parse_failure_is_one_fenced_statement_and_old_lease_cannot_write( tmp_path: Path, ) -> None: repository, connection, _ = _repository_with_one(tmp_path, None) with pytest.raises(LeaseLostError): repository.record_terminal_parse_failure( lease=_job().lease, source=_source(), error_code="INVALID_TEXT_ENCODING", ) statement, parameters = connection.execute.call_args.args assert statement == FAIL_PARSE_SQL assert "FOR UPDATE OF job, document" in statement assert "job.lease_until >= now()" in statement assert "SET status = 'FAILED'" in statement assert "last_error_code = %s" in statement assert parameters[-1] == "INVALID_TEXT_ENCODING" def test_persist_final_fence_loss_raises_inside_transaction_for_rollback( tmp_path: Path, ) -> None: content = b"%PDF-1.7\nsource" source = DocumentSource( upload_id=UPLOAD_ID, document_id=DOCUMENT_ID, knowledge_base_id=KB_ID, access_scope_id=SCOPE_ID, filename="source.pdf", mime_type="application/pdf", storage_key=STORAGE_KEY, byte_size=len(content), raw_sha256=hashlib.sha256(content).hexdigest(), ) artifact = ingest_document( filename=source.filename, declared_mime_type=source.mime_type, content=content, ) plan = plan_artifact(source, artifact, cloud_policy=CloudTextPolicy()) transaction = MagicMock() transaction.__exit__.return_value = False batch_cursor = MagicMock() batch_cursor.__enter__.return_value = batch_cursor connection = MagicMock() connection.__enter__.return_value = connection connection.transaction.return_value = transaction connection.cursor.return_value = batch_cursor def execute(statement: str, parameters: object) -> MagicMock: del parameters cursor = MagicMock() if statement == SELECT_VERSION_SQL: cursor.fetchone.return_value = { "id": plan.version_id, "document_id": DOCUMENT_ID, "parser_profile_hash": plan.parser_profile_hash, "normalization_profile_hash": plan.normalization_profile_hash, "chunk_profile_hash": plan.chunk_profile_hash, "cloud_policy_id": plan.cloud_policy_id, "outbound_manifest_sha256": None, "expected_chunk_count": 0, } elif statement in { SELECT_PAGES_SQL, SELECT_BLOCKS_SQL, SELECT_MANIFEST_ITEMS_SQL, SELECT_CHUNKS_SQL, }: cursor.fetchall.return_value = [] elif statement == UPDATE_DOCUMENT_SQL: cursor.fetchone.return_value = {"id": DOCUMENT_ID} elif statement == FINALIZE_JOB_SQL: cursor.fetchone.return_value = None return cursor connection.execute.side_effect = execute repository = PostgresDocumentWorkflowRepository( _settings(tmp_path), connection_factory=MagicMock(return_value=connection) ) with pytest.raises(LeaseLostError): repository.persist_artifact( lease=_job().lease, source=source, artifact=artifact, cloud_policy=CloudTextPolicy(), embedding_model="text-embedding-v4", embedding_dimension=1024, ) exit_args = transaction.__exit__.call_args.args assert exit_args[0] is LeaseLostError assert "job.lease_until >= now()" in FINALIZE_JOB_SQL assert "job.lease_token = %s" in FINALIZE_JOB_SQL def test_plan_conflicts_fail_closed_before_any_database_write() -> None: content = b"source" source = _source() artifact = ingest_document( filename=source.filename, declared_mime_type=source.mime_type, content=content, ) with pytest.raises(ArtifactConflictError): plan_artifact( source, artifact, cloud_policy=CloudTextPolicy(policy_id="different-policy"), ) def test_jsonb_verification_canonicalizes_nested_tuple_arrays() -> None: value = {"source_anchor": {"block_ids": ("one", "two")}} assert _canonical_json(value) == { "source_anchor": {"block_ids": ["one", "two"]}, }