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()