58 lines
2.5 KiB
Python
58 lines
2.5 KiB
Python
from importlib.util import module_from_spec, spec_from_file_location
|
|
from pathlib import Path
|
|
from urllib.parse import urlsplit
|
|
from uuid import uuid4
|
|
|
|
import pytest
|
|
from alembic.autogenerate import compare_metadata
|
|
from alembic.migration import MigrationContext
|
|
from alembic.operations import Operations
|
|
from sqlalchemy import inspect, text
|
|
|
|
from ai_platform.infrastructure.config.settings import Settings
|
|
from ai_platform.infrastructure.persistence.database import Base, make_session_factory
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
|
|
|
|
async def test_initial_schema_matches_models_and_round_trips():
|
|
settings = Settings.load()
|
|
if urlsplit(settings.database_url).path.rstrip("/").rsplit("/", 1)[-1] != "ai_platform_test":
|
|
pytest.fail("Migration tests require isolated ai_platform_test")
|
|
path = ROOT / "migrations/versions/0001_initial.py"
|
|
spec = spec_from_file_location(path.stem, path)
|
|
assert spec is not None and spec.loader is not None
|
|
migration = module_from_spec(spec)
|
|
spec.loader.exec_module(migration)
|
|
factory = make_session_factory(settings.database_url)
|
|
schema = "migration_test_" + uuid4().hex
|
|
try:
|
|
async with factory.kw["bind"].connect() as connection:
|
|
transaction = await connection.begin()
|
|
try:
|
|
await connection.execute(text(f"CREATE SCHEMA {schema}"))
|
|
await connection.execute(text(f"SET LOCAL search_path TO {schema}"))
|
|
|
|
def apply(sync, direction):
|
|
with Operations.context(MigrationContext.configure(sync)):
|
|
getattr(migration, direction)()
|
|
|
|
def assert_schema(sync):
|
|
context = MigrationContext.configure(
|
|
sync, opts={"compare_server_default": True}
|
|
)
|
|
assert compare_metadata(context, Base.metadata) == []
|
|
assert set(inspect(sync).get_table_names()) == set(Base.metadata.tables)
|
|
|
|
await connection.run_sync(apply, "upgrade")
|
|
await connection.run_sync(assert_schema)
|
|
await connection.run_sync(apply, "downgrade")
|
|
assert await connection.run_sync(lambda sync: inspect(sync).get_table_names()) == []
|
|
await connection.run_sync(apply, "upgrade")
|
|
await connection.run_sync(assert_schema)
|
|
finally:
|
|
# Only this test's schema/data are reverted; public data is untouched.
|
|
await transaction.rollback()
|
|
finally:
|
|
await factory.kw["bind"].dispose()
|