generated from john/python-template
117 lines
4.4 KiB
Python
117 lines
4.4 KiB
Python
"""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 _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)).group_by(Source.document_id)
|
|
)
|
|
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)
|
|
|
|
missing_in_table = sorted(folder_ids - doc_ids)
|
|
missing_in_folders = sorted(doc_ids - folder_ids)
|
|
if missing_in_table:
|
|
folder = missing_in_table[0]
|
|
raise AssertionError(
|
|
"Document directory count and document.doc_id count do not agree. "
|
|
f"./data/documents/{folder} does not appear in document table"
|
|
)
|
|
if missing_in_folders:
|
|
doc_id = missing_in_folders[0]
|
|
raise AssertionError(
|
|
"Document directory count and document.doc_id count do not agree. "
|
|
f"document.doc_id {doc_id} has no corresponding folder in ./data/documents"
|
|
)
|
|
|
|
source_counts = await _source_counts_by_document(settings)
|
|
for document_id in sorted(doc_ids):
|
|
db_count = source_counts.get(document_id, 0)
|
|
file_count = _source_file_count_for_document(upload_dir, document_id)
|
|
if file_count > db_count:
|
|
raise AssertionError(
|
|
"Source file count and source.source_id count do not agree. "
|
|
f"[UPLOAD_DIR]/documents/{document_id}/ contains file(s) with no source row"
|
|
)
|
|
if db_count > file_count:
|
|
raise AssertionError(
|
|
"Source file count and source.source_id count do not agree. "
|
|
f"source.source_id rows exist without files in [UPLOAD_DIR]/documents/{document_id}"
|
|
)
|
|
|
|
|
|
@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="document\\.doc_id count do not agree"):
|
|
await assert_storage_reconciliation(upload_dir=upload_dir, settings=default_settings)
|