harness_zbt/backend/tests/test_migrations.py

import importlib.util
from pathlib import Path
from unittest.mock import Mock

import sqlalchemy as sa
from alembic.config import Config
from alembic.script import ScriptDirectory


def load_core_0004():
    migration_path = (
        Path(__file__).parents[1]
        / "migrations/core/versions/core_0004_product_management.py"
    )
    spec = importlib.util.spec_from_file_location("core_0004_product_management", migration_path)
    assert spec is not None and spec.loader is not None
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def test_each_database_has_the_expected_independent_migration_head() -> None:
    backend_dir = Path(__file__).parents[1]
    expected_heads = {
        "alembic_core": "core_0004",
        "alembic_agent": "agent_0001",
        "alembic_analytics": "analytics_0001",
    }

    for section, expected_head in expected_heads.items():
        config = Config(backend_dir / "alembic.ini", ini_section=section)
        scripts = ScriptDirectory.from_config(config)
        assert scripts.get_current_head() == expected_head


def test_core_0004_skips_columns_that_already_exist(monkeypatch) -> None:
    core_0004_product_management = load_core_0004()
    bind = Mock()
    inspector = Mock()
    inspector.get_columns.return_value = [
        {"name": "premium_cents"},
        {"name": "min_age"},
    ]
    add_column = Mock()
    create_table = Mock()

    monkeypatch.setattr(sa, "inspect", lambda connection: inspector)
    monkeypatch.setattr(core_0004_product_management.op, "get_bind", lambda: bind)
    monkeypatch.setattr(core_0004_product_management.op, "add_column", add_column)
    monkeypatch.setattr(
        core_0004_product_management.CoreBase.metadata.tables["product_change_logs"],
        "create",
        create_table,
    )

    core_0004_product_management.upgrade()

    assert [call.args[1].name for call in add_column.call_args_list] == [
        "coverage_amount_cents",
        "max_age",
    ]
    create_table.assert_called_once_with(bind=bind, checkfirst=True)