transcription/tests/test_storage_reconciliation.py

"""Integrity checks for document/source filesystem-to-database reconciliation."""

from __future__ import annotations

from pathlib import Path
from uuid import UUID
from uuid import uuid4

import pytest
from sqlmodel import func
from sqlmodel import select

from transcription.config import Settings
from transcription.db import session_scope
from transcription.db.models import Document
from transcription.db.models import Source


def _normalize_identifier(value: str) -> str:
    return value.replace("-", "").strip().lower()


def _document_folder_ids(root: Path) -> set[str]:
    documents_root = root / "documents"
    if not documents_root.exists():
        return set()
    return {entry.name for entry in documents_root.iterdir() if entry.is_dir()}


async def _document_ids(settings: Settings) -> set[str]:
    async with session_scope(settings=settings) as session:
        rows = await session.exec(select(Document.id))
        return {str(item) for item in rows.all()}


async def _source_counts_by_document(settings: Settings) -> dict[str, int]:
    async with session_scope(settings=settings) as session:
        rows = await session.exec(
            select(
                Source.document_id,
                func.count(Source.id),  # ty: ignore[invalid-argument-type] - SQLAlchemy descriptor false positive.
            ).group_by(
                Source.document_id  # ty: ignore[invalid-argument-type] - SQLAlchemy descriptor false positive.
            )
        )
        return {str(document_id): int(count) for document_id, count in rows}


def _source_file_count_for_document(root: Path, document_id: str) -> int:
    directory = root / "documents" / document_id
    if not directory.exists():
        return 0
    return sum(1 for entry in directory.iterdir() if entry.is_file())


async def assert_storage_reconciliation(*, upload_dir: Path, settings: Settings) -> None:
    folder_ids = _document_folder_ids(upload_dir)
    doc_ids = await _document_ids(settings)
    folder_by_normalized = {_normalize_identifier(folder_id): folder_id for folder_id in folder_ids}
    doc_by_normalized = {_normalize_identifier(doc_id): doc_id for doc_id in doc_ids}
    source_counts = await _source_counts_by_document(settings)
    source_counts_by_normalized = {
        _normalize_identifier(document_id): count for document_id, count in source_counts.items()
    }
    mismatches: list[str] = []

    missing_in_table = sorted(set(folder_by_normalized) - set(doc_by_normalized))
    for folder_key in missing_in_table:
        folder_name = folder_by_normalized[folder_key]
        mismatches.append(
            f"document-folder-without-row: ./data/documents/{folder_name} has no matching document.doc_id"
        )

    missing_in_folders = sorted(set(doc_by_normalized) - set(folder_by_normalized))
    for doc_key in missing_in_folders:
        doc_id = doc_by_normalized[doc_key]
        source_count = source_counts_by_normalized.get(doc_key, 0)
        mismatches.append(
            "document-row-without-folder: "
            f"document.doc_id {doc_id} has no corresponding folder in ./data/documents "
            f"(source rows: {source_count})"
        )

    for doc_key in sorted(doc_by_normalized):
        doc_id = doc_by_normalized[doc_key]
        folder_name = folder_by_normalized.get(doc_key)
        db_count = source_counts_by_normalized.get(doc_key, 0)
        file_count = _source_file_count_for_document(upload_dir, folder_name) if folder_name is not None else 0
        if db_count != file_count:
            folder_display = (
                f"./data/documents/{folder_name}" if folder_name is not None else "./data/documents/<missing-folder>"
            )
            mismatches.append(
                "source-count-mismatch: "
                f"document.doc_id {doc_id} -> source rows: {db_count}, files in {folder_display}: {file_count}"
            )

    if mismatches:
        report = "\n".join(f"- {item}" for item in mismatches)
        raise AssertionError(
            f"Storage reconciliation mismatch(es) detected.\nReconciling item count: {len(mismatches)}\n{report}"
        )


@pytest.mark.asyncio
async def test_storage_reconciliation_passes_for_matching_counts(tmp_path, default_settings: Settings):
    upload_dir = tmp_path / "uploads"
    document_id = str(uuid4())
    document_uuid = UUID(document_id)
    file_name = f"{uuid4()}.png"
    (upload_dir / "documents" / document_id).mkdir(parents=True, exist_ok=True)
    (upload_dir / "documents" / document_id / file_name).write_bytes(b"ok")

    async with session_scope(settings=default_settings) as session:
        session.add(Document(id=document_uuid, name="Doc"))
        session.add(
            Source(
                document_id=document_uuid,
                page_number=1,
                upload_name=file_name,
                filename=file_name,
                file_path=f"documents/{document_id}/{file_name}",
                file_hash="a" * 64,
                file_size_bytes=2,
            )
        )
        await session.commit()

    await assert_storage_reconciliation(upload_dir=upload_dir, settings=default_settings)


@pytest.mark.asyncio
async def test_storage_reconciliation_reports_actionable_mismatch_message(tmp_path, default_settings: Settings):
    upload_dir = tmp_path / "uploads"
    orphan_dir = upload_dir / "documents" / str(uuid4())
    orphan_dir.mkdir(parents=True, exist_ok=True)

    with pytest.raises(AssertionError, match="Reconciling item count:"):
        await assert_storage_reconciliation(upload_dir=upload_dir, settings=default_settings)


@pytest.mark.asyncio
async def test_storage_reconciliation_reports_missing_document_folder_even_without_sources(
    tmp_path, default_settings: Settings
):
    upload_dir = tmp_path / "uploads"
    missing_folder_doc_id = uuid4()

    async with session_scope(settings=default_settings) as session:
        session.add(Document(id=missing_folder_doc_id, name="Needs folder"))
        await session.commit()

    with pytest.raises(AssertionError, match="document-row-without-folder:"):
        await assert_storage_reconciliation(upload_dir=upload_dir, settings=default_settings)