59 lines
2.2 KiB
Python
59 lines
2.2 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)
|
|
|
|
def test_0040_upgrade_compiles_health_event_indexes(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.20260903_1400_0040_tool_health_indexes"
|
|
)
|
|
with patch.object(migration, "op", operations):
|
|
migration.upgrade()
|
|
|
|
rendered = "\n".join(statements)
|
|
self.assertIn("ix_usage_health_kind_created", rendered)
|
|
self.assertIn("ux_usage_tool_failure_message", rendered)
|
|
self.assertIn("tool_failure", rendered)
|
|
self.assertIn("agent_guard", rendered)
|
|
self.assertIn("context_fold_failure", rendered)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|