97 lines
4.4 KiB
Python
97 lines
4.4 KiB
Python
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
from pathlib import Path
|
|
|
|
import sqlalchemy as sa
|
|
from alembic.migration import MigrationContext
|
|
from alembic.operations import Operations
|
|
|
|
|
|
def _load_migration_module():
|
|
path = (
|
|
Path(__file__).resolve().parents[1]
|
|
/ "alembic"
|
|
/ "versions"
|
|
/ "20260327_add_external_qa_followup_state.py"
|
|
)
|
|
spec = importlib.util.spec_from_file_location("external_qa_followup_state_migration", path)
|
|
module = importlib.util.module_from_spec(spec)
|
|
assert spec is not None
|
|
assert spec.loader is not None
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
def _create_qa_tables(metadata: sa.MetaData) -> None:
|
|
sa.Table(
|
|
"qa_sessions",
|
|
metadata,
|
|
sa.Column("id", sa.Integer(), primary_key=True, autoincrement=True),
|
|
sa.Column("tenant_id", sa.BigInteger(), nullable=False),
|
|
sa.Column("created_at", sa.TIMESTAMP(), nullable=False),
|
|
sa.Column("updated_at", sa.TIMESTAMP(), nullable=False),
|
|
sa.Column("user_id", sa.BigInteger(), nullable=True),
|
|
sa.Column("title", sa.String(length=120), nullable=False),
|
|
sa.Column("project_id", sa.BigInteger(), nullable=True),
|
|
sa.Column("started_at", sa.TIMESTAMP(), nullable=True),
|
|
sa.Column("last_message_at", sa.TIMESTAMP(), nullable=True),
|
|
sa.Column("message_count", sa.Integer(), nullable=False),
|
|
)
|
|
sa.Table(
|
|
"qa_session_states",
|
|
metadata,
|
|
sa.Column("id", sa.Integer(), primary_key=True, autoincrement=True),
|
|
sa.Column("tenant_id", sa.BigInteger(), nullable=False),
|
|
sa.Column("created_at", sa.TIMESTAMP(), nullable=False),
|
|
sa.Column("updated_at", sa.TIMESTAMP(), nullable=False),
|
|
sa.Column("session_id", sa.Integer(), nullable=False),
|
|
sa.Column("last_subject", sa.String(length=255), nullable=True),
|
|
sa.Column("last_entity_ids", sa.JSON(), nullable=True),
|
|
sa.Column("last_intent", sa.String(length=20), nullable=True),
|
|
sa.Column("ambiguity", sa.String(length=20), nullable=True),
|
|
sa.Column("rewrite_history", sa.JSON(), nullable=True),
|
|
)
|
|
|
|
|
|
def test_external_qa_followup_migration_adds_and_removes_columns_and_indexes() -> None:
|
|
migration = _load_migration_module()
|
|
engine = sa.create_engine("sqlite:///:memory:")
|
|
metadata = sa.MetaData()
|
|
_create_qa_tables(metadata)
|
|
metadata.create_all(engine)
|
|
|
|
with engine.begin() as conn:
|
|
context = MigrationContext.configure(conn)
|
|
migration.op = Operations(context)
|
|
migration.upgrade()
|
|
|
|
inspector = sa.inspect(conn)
|
|
session_columns = {column["name"] for column in inspector.get_columns("qa_sessions")}
|
|
state_columns = {column["name"] for column in inspector.get_columns("qa_session_states")}
|
|
session_indexes = {index["name"]: index for index in inspector.get_indexes("qa_sessions")}
|
|
state_indexes = {index["name"]: index for index in inspector.get_indexes("qa_session_states")}
|
|
|
|
assert "external_conversation_key" in session_columns
|
|
assert "pending_clarification" in state_columns
|
|
assert "pending_clarification_expires_at" in state_columns
|
|
assert "ix_qa_sessions_external_conversation_key" in session_indexes
|
|
assert "uq_qa_sessions_tenant_external_conversation_key" in session_indexes
|
|
assert bool(session_indexes["uq_qa_sessions_tenant_external_conversation_key"]["unique"]) is True
|
|
assert "ix_qa_session_states_pending_clarification_expires_at" in state_indexes
|
|
|
|
migration.downgrade()
|
|
|
|
inspector = sa.inspect(conn)
|
|
session_columns = {column["name"] for column in inspector.get_columns("qa_sessions")}
|
|
state_columns = {column["name"] for column in inspector.get_columns("qa_session_states")}
|
|
session_indexes = {index["name"] for index in inspector.get_indexes("qa_sessions")}
|
|
state_indexes = {index["name"] for index in inspector.get_indexes("qa_session_states")}
|
|
|
|
assert "external_conversation_key" not in session_columns
|
|
assert "pending_clarification" not in state_columns
|
|
assert "pending_clarification_expires_at" not in state_columns
|
|
assert "ix_qa_sessions_external_conversation_key" not in session_indexes
|
|
assert "uq_qa_sessions_tenant_external_conversation_key" not in session_indexes
|
|
assert "ix_qa_session_states_pending_clarification_expires_at" not in state_indexes
|