from __future__ import annotations from pathlib import Path import pytest from sqlalchemy import Column, Integer, MetaData, Table, create_engine, text from sqlalchemy.engine import Engine from app.db.migration_preflight import ( LEGACY_ADOPTABLE_HISTORICAL_CASE_TABLES, MIGRATION_OWNED_TABLES_BY_REVISION, MigrationPreflightError, validate_migration_state, ) @pytest.fixture def engine() -> Engine: database = create_engine("sqlite+pysqlite:///:memory:") try: yield database finally: database.dispose() def _create_tables(engine: Engine, table_names: set[str] | frozenset[str]) -> None: metadata = MetaData() for table_name in table_names: Table(table_name, metadata, Column("id", Integer, primary_key=True)) metadata.create_all(engine) def _create_version_table(engine: Engine, *revisions: str) -> None: with engine.begin() as connection: connection.execute(text("CREATE TABLE alembic_version (version_num VARCHAR(32) NOT NULL)")) for revision in revisions: connection.execute( text("INSERT INTO alembic_version (version_num) VALUES (:revision)"), {"revision": revision}, ) def test_unversioned_database_without_migration_owned_tables_is_safe(engine: Engine) -> None: state = validate_migration_state(engine) assert state.revision is None assert state.owned_tables == frozenset() def test_unversioned_database_can_adopt_legacy_historical_case_tables(engine: Engine) -> None: _create_tables(engine, LEGACY_ADOPTABLE_HISTORICAL_CASE_TABLES) state = validate_migration_state(engine) assert state.revision is None assert state.owned_tables == LEGACY_ADOPTABLE_HISTORICAL_CASE_TABLES def test_revision_0007_accepts_partial_legacy_historical_case_tables(engine: Engine) -> None: expected = MIGRATION_OWNED_TABLES_BY_REVISION["20260716_0007"] adopted = frozenset({"risk_observations", "risk_observation_feedback"}) _create_tables(engine, expected | adopted) _create_version_table(engine, "20260716_0007") state = validate_migration_state(engine) assert state.owned_tables == expected | adopted @pytest.mark.parametrize( "owned_table", sorted( MIGRATION_OWNED_TABLES_BY_REVISION["20260716_0023"] - LEGACY_ADOPTABLE_HISTORICAL_CASE_TABLES ), ) def test_unversioned_database_with_any_migration_owned_table_is_rejected( engine: Engine, owned_table: str, ) -> None: _create_tables(engine, {owned_table}) with pytest.raises(MigrationPreflightError, match="unversioned database contains"): validate_migration_state(engine) @pytest.mark.parametrize( ("revision", "expected_tables"), list(MIGRATION_OWNED_TABLES_BY_REVISION.items()), ) def test_known_revision_requires_and_accepts_its_exact_owned_table_set( engine: Engine, revision: str, expected_tables: frozenset[str], ) -> None: _create_tables(engine, expected_tables) _create_version_table(engine, revision) state = validate_migration_state(engine) assert state.revision == revision assert state.owned_tables == expected_tables @pytest.mark.parametrize( ("revision", "actual_tables"), [ ("20260713_0001", frozenset({"expense_cases", "expense_case_links"})), ("20260713_0001", MIGRATION_OWNED_TABLES_BY_REVISION["20260713_0002"]), ( "20260713_0002", MIGRATION_OWNED_TABLES_BY_REVISION["20260713_0002"] - {"auth_sessions"}, ), ( "20260714_0003", MIGRATION_OWNED_TABLES_BY_REVISION["20260714_0003"] - {"ai_decisions"}, ), ( "20260714_0004", MIGRATION_OWNED_TABLES_BY_REVISION["20260714_0004"] - {"ai_application_preview_decisions"}, ), ( "20260714_0005", MIGRATION_OWNED_TABLES_BY_REVISION["20260714_0005"] - {"memory_entries"}, ), ( "20260716_0006", MIGRATION_OWNED_TABLES_BY_REVISION["20260716_0006"] - {"attachment_association_jobs"}, ), ( "20260716_0007", MIGRATION_OWNED_TABLES_BY_REVISION["20260716_0007"] - {"memory_entries"}, ), ( "20260716_0008", MIGRATION_OWNED_TABLES_BY_REVISION["20260716_0008"] - {"few_shot_samples"}, ), ( "20260716_0009", MIGRATION_OWNED_TABLES_BY_REVISION["20260716_0009"] - {"memory_entries"}, ), ( "20260716_0010", MIGRATION_OWNED_TABLES_BY_REVISION["20260716_0010"] - {"approval_action_ledgers"}, ), ( "20260716_0011", MIGRATION_OWNED_TABLES_BY_REVISION["20260716_0011"] - {"risk_disposition_events"}, ), ( "20260716_0012", MIGRATION_OWNED_TABLES_BY_REVISION["20260716_0012"] - {"risk_disposition_events"}, ), ( "20260716_0013", MIGRATION_OWNED_TABLES_BY_REVISION["20260716_0013"] - {"approval_task_events"}, ), ( "20260716_0014", MIGRATION_OWNED_TABLES_BY_REVISION["20260716_0014"] - {"risk_disposition_events"}, ), ( "20260716_0015", MIGRATION_OWNED_TABLES_BY_REVISION["20260716_0015"] - {"savings_events"}, ), ( "20260716_0016", MIGRATION_OWNED_TABLES_BY_REVISION["20260716_0016"] - {"commercial_cost_events"}, ), ( "20260716_0017", MIGRATION_OWNED_TABLES_BY_REVISION["20260716_0017"] - {"financial_connector_events"}, ), ( "20260716_0018", MIGRATION_OWNED_TABLES_BY_REVISION["20260716_0018"] - {"agent_asset_release_observations"}, ), ( "20260716_0019", MIGRATION_OWNED_TABLES_BY_REVISION["20260716_0019"] - {"commercial_runtime_reservations"}, ), ( "20260716_0020", MIGRATION_OWNED_TABLES_BY_REVISION["20260716_0020"] - {"financial_connector_config_events"}, ), ( "20260716_0021", MIGRATION_OWNED_TABLES_BY_REVISION["20260716_0021"] - {"commercial_billing_periods"}, ), ( "20260716_0022", MIGRATION_OWNED_TABLES_BY_REVISION["20260716_0022"] - {"financial_connector_operational_events"}, ), ( "20260716_0023", MIGRATION_OWNED_TABLES_BY_REVISION["20260716_0023"] - {"agent_asset_release_audit_samples"}, ), ( "20260717_0025", MIGRATION_OWNED_TABLES_BY_REVISION["20260717_0025"] - {"tenants"}, ), ( "20260717_0027", MIGRATION_OWNED_TABLES_BY_REVISION["20260717_0027"] - {"knowledge_onlyoffice_sessions"}, ), ( "20260718_0029", MIGRATION_OWNED_TABLES_BY_REVISION["20260718_0029"] - {"tenant_finance_report_runs"}, ), ], ) def test_known_revision_with_missing_or_unexpected_owned_tables_is_rejected( engine: Engine, revision: str, actual_tables: frozenset[str], ) -> None: _create_tables(engine, actual_tables) _create_version_table(engine, revision) with pytest.raises(MigrationPreflightError, match="does not match revision"): validate_migration_state(engine) def test_unknown_revision_is_rejected(engine: Engine) -> None: _create_version_table(engine, "20990101_unknown") with pytest.raises(MigrationPreflightError, match="unknown Alembic revision"): validate_migration_state(engine) def test_multiple_revisions_are_rejected(engine: Engine) -> None: _create_version_table(engine, "20260713_0001", "20260713_0002") with pytest.raises(MigrationPreflightError, match="multiple Alembic revisions"): validate_migration_state(engine) def test_empty_version_table_is_safe_only_when_owned_tables_are_absent(engine: Engine) -> None: _create_version_table(engine) assert validate_migration_state(engine).revision is None _create_tables(engine, {"expense_cases"}) with pytest.raises(MigrationPreflightError, match="no recorded revision"): validate_migration_state(engine) def test_server_start_runs_preflight_before_alembic_upgrade() -> None: script_path = Path(__file__).resolve().parents[1] / "server_start.sh" script = script_path.read_text(encoding="utf-8") preflight = '"$PYTHON_BIN" -m app.db.migration_preflight' upgrade = '"$PYTHON_BIN" -m alembic -c "$SCRIPT_DIR/alembic.ini" upgrade head' assert 'PYTHONPATH="$SCRIPT_DIR/src${PYTHONPATH:+:$PYTHONPATH}"' in script assert script.index(preflight) < script.index(upgrade)