Function test_automigration

tests/units/test_sqlalchemy.py:146–296

Test alembic automigration with add and drop table and column. Args: tmp_working_dir: directory where database and migrations are stored monkeypatch: pytest fixture to overwrite attributes model_registry: clean reflex ModelRegistry

(
    tmp_working_dir: Path,
    monkeypatch: pytest.MonkeyPatch,
    model_registry: type[ModelRegistry],
)

    "ignore:This declarative base already contains a class with the same class name",
)
def test_automigration(
    tmp_working_dir: Path,
    monkeypatch: pytest.MonkeyPatch,
    model_registry: type[ModelRegistry],
):
    """Test alembic automigration with add and drop table and column.

    Args:
        tmp_working_dir: directory where database and migrations are stored
        monkeypatch: pytest fixture to overwrite attributes
        model_registry: clean reflex ModelRegistry
    """
    from sqlalchemy import select
    from sqlalchemy.exc import OperationalError
    from sqlalchemy.orm import (
        DeclarativeBase,
        Mapped,
        MappedAsDataclass,
        declared_attr,
        mapped_column,
    )

    alembic_ini = tmp_working_dir / "alembic.ini"
    versions = tmp_working_dir / "alembic" / "versions"
    monkeypatch.setattr(reflex.constants, "ALEMBIC_CONFIG", str(alembic_ini))

    config_mock = mock.Mock()
    config_mock.db_url = f"sqlite:///{tmp_working_dir}/reflex.db"
    monkeypatch.setattr(reflex.model, "get_config", mock.Mock(return_value=config_mock))

    assert alembic_ini.exists() is False
    assert versions.exists() is False
    alembic_init()
    assert alembic_ini.exists()
    assert versions.exists()

    class Base(DeclarativeBase):
        @declared_attr.directive
        def __tablename__(cls) -> str:
            return cls.__name__.lower()

    assert model_registry.register(Base)

    class ModelBase(Base, MappedAsDataclass):
        __abstract__ = True
        id: Mapped[int | None] = mapped_column(primary_key=True, default=None)

    # initial table
    class AlembicThing(ModelBase):  # pyright: ignore[reportRedeclaration]
        t1: Mapped[str] = mapped_column(default="")

    with get_engine().connect() as connection:
        assert alembic_autogenerate(connection=connection, message="Initial Revision")
    assert migrate()
    version_scripts = list(versions.glob("*.py"))
    assert len(version_scripts) == 1
    assert version_scripts[0].name.endswith("initial_revision.py")