zcbot/tests/test_storage_migration.py

38 lines
1.3 KiB
Python

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