from __future__ import annotations import tempfile import unittest from pathlib import Path from alembic import command from alembic.runtime.migration import MigrationContext from alembic.script import ScriptDirectory from sqlalchemy import create_engine, inspect, text from govoplan_core.core.change_sequence import ChangeSequenceEntry from govoplan_core.db.migrations import ( alembic_config, migrate_database, reconcile_change_sequence_retention_floor_drift, reconcile_namespace_table_drift, ) def configured_migration_heads( database_url: str, *, enabled_modules: tuple[str, ...] | list[str] | None = None, migration_track: str = "release", ) -> set[str]: script = ScriptDirectory.from_config( alembic_config( database_url=database_url, enabled_modules=enabled_modules, migration_track=migration_track, ), ) return {revision.revision for revision in script.get_revisions("heads")} def database_migration_heads(connection) -> set[str]: return set(MigrationContext.configure(connection).get_current_heads()) class DatabaseMigrationTests(unittest.TestCase): def test_migration_tracks_use_separate_version_locations(self) -> None: release_locations = alembic_config(database_url="sqlite:////tmp/govoplan-release.db").get_main_option("version_locations") dev_locations = alembic_config( database_url="sqlite:////tmp/govoplan-dev.db", migration_track="dev", ).get_main_option("version_locations") self.assertIn("alembic/versions", release_locations) self.assertNotIn("alembic/dev_versions", release_locations) self.assertIn("alembic/dev_versions", dev_locations) self.assertIn("migrations/dev_versions", dev_locations) def test_repairs_missing_change_sequence_retention_floor(self) -> None: with tempfile.TemporaryDirectory(prefix="govoplan-change-sequence-drift-test-") as directory: database = Path(directory) / "change-sequence-drift.db" url = f"sqlite:///{database}" engine = create_engine(url) try: with engine.begin() as connection: ChangeSequenceEntry.__table__.create(bind=connection) with engine.connect() as connection: tables = set(inspect(connection).get_table_names()) self.assertIn("core_change_sequence", tables) self.assertNotIn("core_change_sequence_retention_floor", tables) finally: engine.dispose() self.assertTrue(reconcile_change_sequence_retention_floor_drift(url)) engine = create_engine(url) try: with engine.connect() as connection: inspector = inspect(connection) self.assertIn("core_change_sequence_retention_floor", inspector.get_table_names()) self.assertIn( "ix_core_change_sequence_retention_scope", {index["name"] for index in inspector.get_indexes("core_change_sequence_retention_floor")}, ) finally: engine.dispose() def test_repairs_namespace_tables_when_database_was_stamped_ahead(self) -> None: with tempfile.TemporaryDirectory(prefix="govoplan-namespace-drift-test-") as directory: database = Path(directory) / "namespace-drift.db" url = f"sqlite:///{database}" engine = create_engine(url) try: with engine.begin() as connection: connection.execute(text("CREATE TABLE accounts (id TEXT PRIMARY KEY)")) connection.execute(text("CREATE TABLE users (id TEXT PRIMARY KEY)")) connection.execute(text("CREATE TABLE governance_templates (id TEXT PRIMARY KEY)")) with engine.connect() as connection: tables = set(inspect(connection).get_table_names()) self.assertIn("accounts", tables) self.assertNotIn("access_accounts", tables) finally: engine.dispose() self.assertTrue(reconcile_namespace_table_drift(url)) engine = create_engine(url) try: with engine.connect() as connection: tables = set(inspect(connection).get_table_names()) self.assertNotIn("accounts", tables) self.assertIn("access_accounts", tables) self.assertNotIn("governance_templates", tables) self.assertIn("admin_governance_templates", tables) finally: engine.dispose() def test_core_and_access_baselines_apply_to_fresh_database(self) -> None: with tempfile.TemporaryDirectory(prefix="govoplan-core-baseline-test-") as directory: database = Path(directory) / "core.db" url = f"sqlite:///{database}" command.upgrade(alembic_config(database_url=url, enabled_modules=()), "heads") engine = create_engine(url) try: with engine.connect() as connection: inspector = inspect(connection) tables = set(inspector.get_table_names()) role_flags = { row["slug"]: bool(row["is_builtin"]) for row in connection.execute( text( """ SELECT slug, is_builtin FROM access_roles WHERE tenant_id IS NULL AND slug IN ('system_owner', 'system_admin', 'system_auditor') """ ), ).mappings() } system_settings_count = connection.execute( text("SELECT COUNT(*) FROM core_system_settings WHERE id = 'global'"), ).scalar_one() current = database_migration_heads(connection) self.assertEqual(current, configured_migration_heads(url, enabled_modules=())) self.assertIn("core_scopes", tables) self.assertNotIn("tenancy_tenants", tables) self.assertIn("access_accounts", tables) self.assertIn("access_users", tables) self.assertIn("access_identities", tables) self.assertIn("access_function_assignments", tables) self.assertIn("audit_log", tables) self.assertIn("audit_outbox_events", tables) self.assertIn("file_assets", tables) self.assertIn("mail_server_profiles", tables) self.assertEqual(system_settings_count, 1) self.assertEqual( role_flags, {"system_owner": True, "system_admin": False, "system_auditor": False}, ) finally: engine.dispose() def test_default_module_baselines_apply_to_fresh_database(self) -> None: with tempfile.TemporaryDirectory(prefix="govoplan-default-baseline-test-") as directory: database = Path(directory) / "default.db" url = f"sqlite:///{database}" result = migrate_database(database_url=url) engine = create_engine(url) try: with engine.connect() as connection: inspector = inspect(connection) tables = set(inspector.get_table_names()) current = database_migration_heads(connection) mail_policy_fks = inspector.get_foreign_keys("mail_profile_policies") mail_referred_tables = {fk["referred_table"] for fk in mail_policy_fks} folder_indexes = {index["name"] for index in inspector.get_indexes("file_folders")} import_profile_indexes = { index["name"] for index in inspector.get_indexes("campaign_recipient_import_mapping_profiles") } self.assertIsNone(result.previous_revision) self.assertIsNone(result.reconciled_revision) self.assertEqual(current, configured_migration_heads(url)) self.assertEqual(result.current_revision, ",".join(sorted(current))) self.assertIn("calendar_sync_credentials", tables) self.assertIn("campaign_recipient_import_mapping_profiles", tables) self.assertIn("file_connector_credentials", tables) self.assertIn("file_connector_policies", tables) self.assertIn("identity_identities", tables) self.assertIn("identity_account_links", tables) self.assertIn("mail_profile_policies", tables) self.assertIn("organizations_functions", tables) self.assertIn("core_scopes", mail_referred_tables) self.assertIn("uq_file_folders_active_user_path", folder_indexes) self.assertIn("uq_file_folders_active_group_path", folder_indexes) self.assertIn("ix_recipient_import_profiles_ordered_fp", import_profile_indexes) finally: engine.dispose() def test_dev_migration_track_applies_detailed_chain_to_fresh_database(self) -> None: with tempfile.TemporaryDirectory(prefix="govoplan-dev-track-migration-test-") as directory: database = Path(directory) / "dev-track.db" url = f"sqlite:///{database}" result = migrate_database(database_url=url, migration_track="dev") engine = create_engine(url) try: with engine.connect() as connection: tables = set(inspect(connection).get_table_names()) current = database_migration_heads(connection) self.assertEqual(current, configured_migration_heads(url, migration_track="dev")) self.assertEqual(result.current_revision, ",".join(sorted(current))) self.assertIn("0f1e2d3c4b5a", current) self.assertIn("audit_outbox_events", tables) self.assertIn("file_connector_profiles", tables) self.assertIn("core_scopes", tables) finally: engine.dispose() def test_migrate_database_removes_stale_dependency_parent_version_rows(self) -> None: with tempfile.TemporaryDirectory(prefix="govoplan-stale-head-baseline-test-") as directory: database = Path(directory) / "stale-heads.db" url = f"sqlite:///{database}" migrate_database(database_url=url) engine = create_engine(url) try: with engine.begin() as connection: connection.execute( text("INSERT INTO alembic_version (version_num) VALUES (:revision)"), {"revision": "4f2a9c8e7b6d"}, ) with engine.connect() as connection: self.assertIn("4f2a9c8e7b6d", database_migration_heads(connection)) finally: engine.dispose() result = migrate_database(database_url=url) engine = create_engine(url) try: with engine.connect() as connection: current = database_migration_heads(connection) self.assertEqual(current, configured_migration_heads(url)) self.assertNotIn("4f2a9c8e7b6d", current) self.assertEqual(result.current_revision, ",".join(sorted(current))) finally: engine.dispose() def test_full_product_baselines_apply_with_idm(self) -> None: enabled_modules = [ "tenancy", "organizations", "identity", "access", "admin", "dashboard", "policy", "audit", "files", "mail", "campaigns", "calendar", "docs", "ops", "idm", ] with tempfile.TemporaryDirectory(prefix="govoplan-full-baseline-test-") as directory: database = Path(directory) / "full.db" url = f"sqlite:///{database}" result = migrate_database(database_url=url, enabled_modules=enabled_modules) engine = create_engine(url) try: with engine.connect() as connection: tables = set(inspect(connection).get_table_names()) current = database_migration_heads(connection) self.assertEqual(current, configured_migration_heads(url, enabled_modules=enabled_modules)) self.assertEqual(result.current_revision, ",".join(sorted(current))) self.assertIn("idm_tenant_settings", tables) self.assertIn("idm_organization_function_assignments", tables) finally: engine.dispose() if __name__ == "__main__": unittest.main()