from __future__ import annotations import importlib import unittest from unittest.mock import patch from alembic.migration import MigrationContext from alembic.operations import Operations from sqlalchemy import create_mock_engine from sqlalchemy.dialects import postgresql class StorageMigrationTests(unittest.TestCase): def test_0029_upgrade_compiles_as_postgresql_ddl(self) -> None: statements: list[str] = [] def capture(sql, *multiparams, **params): statements.append(str(sql.compile(dialect=postgresql.dialect()))) engine = create_mock_engine("postgresql+psycopg://", capture) operations = Operations(MigrationContext.configure(engine.connect())) migration = importlib.import_module( "db.migrations.versions.20260812_1800_0029_storage_hot_paths" ) with patch.object(migration, "op", operations): migration.upgrade() rendered = "\n".join(statements) self.assertIn("next_message_idx", rendered) self.assertIn("MAX(idx) + 1", rendered) self.assertIn("ix_tasks_user_active_updated", rendered) self.assertIn("ix_scheduled_jobs_due_active", rendered) self.assertIn("USING brin", rendered) if __name__ == "__main__": unittest.main()