X-Financial/server/tests/test_migration_preflight.py

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)