govoplan-files/tests/test_integrity_reconciliation.py

from __future__ import annotations

import hashlib
import tempfile
import unittest
from pathlib import Path

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

from govoplan_access.backend.db.models import Account, Group, User
from govoplan_core.db.base import Base
from govoplan_files.backend.db.models import (
    FileBlob,
    FileIntegrityFinding,
    FileIntegrityScan,
)
from govoplan_files.backend.storage.backends import LocalFilesystemStorageBackend
from govoplan_files.backend.storage.common import FileStorageError
from govoplan_files.backend.storage.integrity import (
    cleanup_orphan_finding,
    create_integrity_scan,
    read_verified_blob_bytes,
    recheck_integrity_finding,
    run_integrity_scan_batch,
)


TENANT_ID = "tenant-1"
USER_ID = "user-1"


class IntegrityReconciliationTests(unittest.TestCase):
    def setUp(self) -> None:
        self.temporary_directory = tempfile.TemporaryDirectory()
        self.addCleanup(self.temporary_directory.cleanup)
        self.backend = LocalFilesystemStorageBackend(
            Path(self.temporary_directory.name)
        )
        self.engine = create_engine("sqlite:///:memory:", future=True)
        Base.metadata.create_all(
            bind=self.engine,
            tables=[
                Account.__table__,
                User.__table__,
                Group.__table__,
                FileBlob.__table__,
                FileIntegrityScan.__table__,
                FileIntegrityFinding.__table__,
            ],
        )
        self.session = sessionmaker(bind=self.engine, future=True)()
        self.addCleanup(self._close)

    def _close(self) -> None:
        self.session.close()
        self.engine.dispose()

    def test_scan_is_bounded_resumable_and_reconciles_both_orphan_directions(
        self,
    ) -> None:
        valid_data = b"valid"
        restored_data = b"restore-me"
        expected_corrupt_data = b"expected"
        stored_corrupt_data = b"corrupt!"
        valid = _blob("blob-1", "valid.bin", valid_data)
        missing = _blob("blob-2", "missing.bin", restored_data)
        corrupt = _blob("blob-3", "corrupt.bin", expected_corrupt_data)
        self.session.add_all([valid, missing, corrupt])
        self.session.commit()
        self.backend.put_bytes(valid.storage_key, valid_data)
        self.backend.put_bytes(corrupt.storage_key, stored_corrupt_data)
        orphan_key = f"tenants/{TENANT_ID}/files/orphan.bin"
        self.backend.put_bytes(orphan_key, b"orphan")

        scan = create_integrity_scan(
            self.session,
            tenant_id=TENANT_ID,
            user_id=USER_ID,
            batch_size=1,
            backend=self.backend,
        )
        self.session.commit()

        invocations = 0
        while scan.status != "completed":
            run_integrity_scan_batch(
                self.session,
                scan,
                backend=self.backend,
            )
            self.session.commit()
            invocations += 1
            scan = self.session.get(FileIntegrityScan, scan.id)
            self.assertIsNotNone(scan)
            self.session.expire_all()
            if invocations > 12:
                self.fail("Integrity scan did not complete")

        self.assertGreater(invocations, 3)
        self.assertEqual(3, scan.scanned_blob_count)
        self.assertEqual(1, scan.verified_blob_count)
        self.assertEqual(2, scan.quarantined_blob_count)
        self.assertEqual(3, scan.scanned_object_count)
        self.assertEqual(1, scan.orphan_object_count)
        findings = (
            self.session.query(FileIntegrityFinding)
            .filter(FileIntegrityFinding.scan_id == scan.id)
            .all()
        )
        self.assertEqual(
            {"missing", "checksum_mismatch", "orphan_object"},
            {finding.kind for finding in findings},
        )

        missing = self.session.get(FileBlob, missing.id)
        corrupt = self.session.get(FileBlob, corrupt.id)
        self.assertEqual("missing", missing.integrity_status)
        self.assertEqual("checksum_mismatch", corrupt.integrity_status)
        with self.assertRaisesRegex(FileStorageError, "quarantined"):
            read_verified_blob_bytes(missing, backend=self.backend)

        missing_finding = next(
            finding for finding in findings if finding.kind == "missing"
        )
        self.backend.put_bytes(missing.storage_key, restored_data)
        preview = recheck_integrity_finding(
            self.session,
            missing_finding,
            user_id=USER_ID,
            dry_run=True,
            backend=self.backend,
        )
        self.assertTrue(preview.inspection.valid)
        self.assertEqual("open", missing_finding.state)
        repaired = recheck_integrity_finding(
            self.session,
            missing_finding,
            user_id=USER_ID,
            dry_run=False,
            backend=self.backend,
        )
        self.session.commit()
        self.assertTrue(repaired.changed)
        self.assertEqual("resolved", missing_finding.state)
        self.assertEqual(
            restored_data,
            read_verified_blob_bytes(missing, backend=self.backend),
        )

        orphan_finding = next(
            finding for finding in findings if finding.kind == "orphan_object"
        )
        preview_cleanup = cleanup_orphan_finding(
            self.session,
            orphan_finding,
            user_id=USER_ID,
            dry_run=True,
            backend=self.backend,
        )
        self.assertEqual("would_delete", preview_cleanup.action)
        self.assertTrue(self.backend.exists(orphan_key))
        cleanup = cleanup_orphan_finding(
            self.session,
            orphan_finding,
            user_id=USER_ID,
            dry_run=False,
            backend=self.backend,
        )
        repeated = cleanup_orphan_finding(
            self.session,
            orphan_finding,
            user_id=USER_ID,
            dry_run=False,
            backend=self.backend,
        )
        self.assertEqual("deleted", cleanup.action)
        self.assertFalse(self.backend.exists(orphan_key))
        self.assertEqual("already_deleted", repeated.action)
        self.assertFalse(repeated.changed)


def _blob(blob_id: str, filename: str, expected_data: bytes) -> FileBlob:
    return FileBlob(
        id=blob_id,
        tenant_id=TENANT_ID,
        storage_backend="local",
        storage_key=f"tenants/{TENANT_ID}/files/{filename}",
        checksum_sha256=hashlib.sha256(expected_data).hexdigest(),
        size_bytes=len(expected_data),
        ref_count=1,
    )


if __name__ == "__main__":
    unittest.main()