RAG/backend/tests/unit/test_document_workflows.py

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"]},
    }