wehub-resource-sync
85 lines
2.5 KiB
Python
85 lines
2.5 KiB
Python
"""Tests for the conversation_items PK-widening migration (y1a2b3c4d5e6)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
|
|
import sqlalchemy as sa
|
|
from alembic import command
|
|
from sqlalchemy.engine import Engine
|
|
|
|
from omnigent.db.utils import (
|
|
_build_alembic_config,
|
|
clear_engine_cache,
|
|
get_or_create_engine,
|
|
)
|
|
|
|
_TABLE = "conversation_items"
|
|
_HEAD_PK = ["workspace_id", "conversation_id", "id"]
|
|
_PRIOR_PK = ["workspace_id", "id"]
|
|
|
|
|
|
def _pk(engine: Engine) -> list[str]:
|
|
return sa.inspect(engine).get_pk_constraint(_TABLE)["constrained_columns"]
|
|
|
|
|
|
def test_head_widens_conversation_items_pk(tmp_path: Path) -> None:
|
|
"""At head the PK is ``(workspace_id, conversation_id, id)``."""
|
|
uri = f"sqlite:///{tmp_path / 'head.db'}"
|
|
engine = get_or_create_engine(uri)
|
|
try:
|
|
assert _pk(engine) == _HEAD_PK
|
|
finally:
|
|
engine.dispose()
|
|
clear_engine_cache()
|
|
|
|
|
|
def test_downgrade_restores_prior_conversation_items_pk(tmp_path: Path) -> None:
|
|
"""Downgrading one step drops conversation_id back out of the PK."""
|
|
uri = f"sqlite:///{tmp_path / 'downgrade.db'}"
|
|
engine = get_or_create_engine(uri)
|
|
try:
|
|
assert _pk(engine) == _HEAD_PK
|
|
|
|
config = _build_alembic_config(uri)
|
|
with engine.begin() as conn:
|
|
config.attributes["connection"] = conn
|
|
command.downgrade(config, "x1a2b3c4d5e6")
|
|
|
|
assert _pk(engine) == _PRIOR_PK
|
|
finally:
|
|
engine.dispose()
|
|
clear_engine_cache()
|
|
|
|
|
|
def test_items_round_trip_via_store(tmp_path: Path) -> None:
|
|
"""Append + list still works with conversation_id in the PK."""
|
|
from omnigent.entities.conversation import MessageData, NewConversationItem
|
|
from omnigent.stores.conversation_store.sqlalchemy_store import (
|
|
SqlAlchemyConversationStore,
|
|
)
|
|
|
|
uri = f"sqlite:///{tmp_path / 'roundtrip.db'}"
|
|
engine = get_or_create_engine(uri)
|
|
try:
|
|
store = SqlAlchemyConversationStore(str(engine.url))
|
|
conv = store.create_conversation()
|
|
store.append(
|
|
conv.id,
|
|
[
|
|
NewConversationItem(
|
|
type="message",
|
|
response_id="resp_1",
|
|
data=MessageData(
|
|
role="user",
|
|
content=[{"type": "input_text", "text": "hi"}],
|
|
),
|
|
)
|
|
],
|
|
)
|
|
items = store.list_items(conv.id).data
|
|
assert [i.type for i in items] == ["message"]
|
|
finally:
|
|
engine.dispose()
|
|
clear_engine_cache()
|