from __future__ import annotations import tempfile import unittest from pathlib import Path from alembic.runtime.migration import MigrationContext from sqlalchemy import create_engine, inspect from govoplan_core.db.migrations import migrate_database from govoplan_encryption.backend.manifest import get_manifest class EncryptionMigrationTests(unittest.TestCase): def test_migration_creates_lifecycle_tables(self) -> None: with tempfile.TemporaryDirectory(prefix="govoplan-encryption-") as directory: url = f"sqlite:///{Path(directory) / 'encryption.db'}" migrate_database( database_url=url, enabled_modules=("encryption",), manifest_factories=(get_manifest,), ) engine = create_engine(url) try: tables = set(inspect(engine).get_table_names()) self.assertTrue( { "encryption_vaults", "encryption_key_versions", "encryption_key_operations", "encryption_content_protections", "encryption_protection_migrations", "encryption_recovery_ceremonies", "encryption_recovery_approvals", }.issubset(tables) ) with engine.connect() as connection: self.assertIn( "d4a6b8c0e2f3", set(MigrationContext.configure(connection).get_current_heads()), ) finally: engine.dispose() if __name__ == "__main__": unittest.main()