45 lines
1.7 KiB
Python
45 lines
1.7 KiB
Python
from __future__ import annotations
|
|
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
from alembic import command
|
|
from alembic.config import Config
|
|
from sqlalchemy import inspect, select
|
|
|
|
from services.control_api.persistence.db import create_engine_from_url, session_factory
|
|
from services.control_api.persistence.models import Base, WorkspaceRow
|
|
|
|
|
|
class MigrationTests(unittest.TestCase):
|
|
def _config(self, database_path: Path) -> Config:
|
|
config = Config("infrastructure/database/alembic.ini")
|
|
config.set_main_option("sqlalchemy.url", f"sqlite:///{database_path.as_posix()}")
|
|
return config
|
|
|
|
def test_empty_upgrade_downgrade_and_reupgrade(self) -> None:
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
path = Path(directory) / "migration.db"
|
|
config = self._config(path)
|
|
command.upgrade(config, "head")
|
|
engine = create_engine_from_url(f"sqlite:///{path}")
|
|
self.assertIn("workspaces", inspect(engine).get_table_names())
|
|
self.assertIn(
|
|
"retention_until",
|
|
{column["name"] for column in inspect(engine).get_columns("audit_events")},
|
|
)
|
|
self.assertEqual(len(Base.metadata.tables), 20)
|
|
command.downgrade(config, "base")
|
|
self.assertNotIn("workspaces", inspect(engine).get_table_names())
|
|
engine.dispose()
|
|
command.upgrade(config, "head")
|
|
engine = create_engine_from_url(f"sqlite:///{path}")
|
|
with session_factory(engine)() as session:
|
|
self.assertEqual(session.scalars(select(WorkspaceRow)).all(), [])
|
|
engine.dispose()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|