65 lines
2.6 KiB
Python
65 lines
2.6 KiB
Python
from __future__ import annotations
|
|
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
from alembic import command
|
|
from alembic.runtime.migration import MigrationContext
|
|
from sqlalchemy import create_engine, inspect
|
|
|
|
from app.db.base import Base
|
|
from app.db.migrations import (
|
|
REVISION_AUTH_RBAC,
|
|
REVISION_FILE_FOLDERS,
|
|
alembic_config,
|
|
migrate_database,
|
|
)
|
|
|
|
|
|
class DatabaseMigrationTests(unittest.TestCase):
|
|
def test_repairs_create_all_schema_drift_and_upgrades_to_head(self) -> None:
|
|
with tempfile.TemporaryDirectory(prefix="msm-migration-test-") as directory:
|
|
database = Path(directory) / "legacy.db"
|
|
url = f"sqlite:///{database}"
|
|
|
|
# Reproduce the historical development database: Alembic was run
|
|
# through auth/RBAC, then create_all() created later file tables
|
|
# without advancing alembic_version and without altering the
|
|
# already-existing campaign_versions table.
|
|
command.upgrade(alembic_config(database_url=url), REVISION_AUTH_RBAC)
|
|
engine = create_engine(url)
|
|
try:
|
|
Base.metadata.create_all(bind=engine)
|
|
with engine.connect() as connection:
|
|
self.assertEqual(
|
|
MigrationContext.configure(connection).get_current_revision(),
|
|
REVISION_AUTH_RBAC,
|
|
)
|
|
self.assertIn("file_blobs", inspect(connection).get_table_names())
|
|
self.assertNotIn(
|
|
"user_lock_state",
|
|
{column["name"] for column in inspect(connection).get_columns("campaign_versions")},
|
|
)
|
|
finally:
|
|
engine.dispose()
|
|
|
|
result = migrate_database(database_url=url)
|
|
self.assertEqual(result.previous_revision, REVISION_AUTH_RBAC)
|
|
self.assertEqual(result.reconciled_revision, REVISION_FILE_FOLDERS)
|
|
|
|
engine = create_engine(url)
|
|
try:
|
|
with engine.connect() as connection:
|
|
current = MigrationContext.configure(connection).get_current_revision()
|
|
columns = {
|
|
column["name"]
|
|
for column in inspect(connection).get_columns("campaign_versions")
|
|
}
|
|
self.assertEqual(current, result.current_revision)
|
|
self.assertIn("user_lock_state", columns)
|
|
self.assertIn("user_locked_at", columns)
|
|
self.assertIn("user_locked_by_user_id", columns)
|
|
finally:
|
|
engine.dispose()
|