157 lines
6.4 KiB
Python
157 lines
6.4 KiB
Python
"""Exercise real withdrawal SQL and guard against loading an entire long conversation."""
|
|
|
|
from datetime import UTC, datetime, timedelta
|
|
|
|
import pytest
|
|
from app.core.exceptions import AppException
|
|
from app.models.entities import (
|
|
Conversation,
|
|
ConversationMessage,
|
|
IntegrationAssessment,
|
|
IntegrationConversation,
|
|
OperationAudit,
|
|
UserAccount,
|
|
)
|
|
from app.services.conversations import ConversationService
|
|
from sqlalchemy import JSON, BigInteger, Integer, MetaData, event, inspect, select
|
|
from sqlalchemy.dialects.postgresql import JSONB
|
|
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
|
|
|
|
|
@pytest.fixture
|
|
async def withdraw_db(tmp_path):
|
|
engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'withdraw.db'}")
|
|
schema = MetaData()
|
|
for model in (
|
|
UserAccount,
|
|
Conversation,
|
|
ConversationMessage,
|
|
OperationAudit,
|
|
IntegrationAssessment,
|
|
IntegrationConversation,
|
|
):
|
|
copied = model.__table__.to_metadata(schema)
|
|
for column in copied.columns:
|
|
if isinstance(column.type, JSONB):
|
|
column.type = JSON()
|
|
elif isinstance(column.type, BigInteger):
|
|
column.type = Integer()
|
|
async with engine.begin() as connection:
|
|
await connection.run_sync(schema.create_all)
|
|
sessions = async_sessionmaker(engine, expire_on_commit=False)
|
|
start = datetime(2026, 10, 27, tzinfo=UTC)
|
|
async with sessions() as session:
|
|
session.add_all([UserAccount(id=i, emp_no=f"u{i}", name=f"用户{i}", role="user") for i in (1, 2)])
|
|
session.add_all([Conversation(id=i, user_id=i, title=f"会话{i}") for i in (1, 2)])
|
|
for message_id, role, minute in (
|
|
(1, "user", 0),
|
|
(2, "assistant", 0),
|
|
(3, "user", 1),
|
|
(4, "user", 2),
|
|
(5, "assistant", 2),
|
|
(6, "user", 3),
|
|
):
|
|
session.add(
|
|
ConversationMessage(
|
|
id=message_id,
|
|
conversation_id=1,
|
|
role=role,
|
|
content=f"消息{message_id}",
|
|
created_at=start + timedelta(minutes=minute),
|
|
)
|
|
)
|
|
session.add(ConversationMessage(id=7, conversation_id=2, role="user", content="他人的消息"))
|
|
await session.commit()
|
|
try:
|
|
yield sessions, engine, start
|
|
finally:
|
|
await engine.dispose()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(("target_id", "deleted"), [(1, [1, 2]), (3, [3]), (6, [6])])
|
|
async def test_withdraw_only_removes_target_and_immediate_assistant(withdraw_db, target_id, deleted):
|
|
sessions, _, _ = withdraw_db
|
|
async with sessions() as session:
|
|
user = await session.get(UserAccount, 1)
|
|
result = await ConversationService(session).withdraw_turn(1, target_id, user)
|
|
assert result.messages_total == 6 - len(deleted)
|
|
assert [message.id for message in result.messages] == [i for i in range(1, 7) if i not in deleted]
|
|
await session.commit()
|
|
async with sessions() as session:
|
|
ids = list(await session.scalars(select(ConversationMessage.id).order_by(ConversationMessage.id)))
|
|
assert ids == [i for i in range(1, 8) if i not in deleted]
|
|
audit = await session.scalar(select(OperationAudit))
|
|
assert audit.action == "conversation.withdraw"
|
|
assert audit.audit_metadata == {"deleted_message_ids": deleted}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
("conversation_id", "message_id", "user_id", "status"),
|
|
[(1, 1, 2, 404), (999, 1, 1, 404), (1, 999, 1, 422), (1, 2, 1, 422), (1, 7, 1, 422)],
|
|
)
|
|
async def test_withdraw_rejects_other_users_and_invalid_targets(
|
|
withdraw_db, conversation_id, message_id, user_id, status
|
|
):
|
|
sessions, _, _ = withdraw_db
|
|
async with sessions() as session:
|
|
user = await session.get(UserAccount, user_id)
|
|
with pytest.raises(AppException) as error:
|
|
await ConversationService(session).withdraw_turn(conversation_id, message_id, user)
|
|
assert error.value.status_code == status
|
|
await session.rollback()
|
|
async with sessions() as session:
|
|
assert len(list(await session.scalars(select(ConversationMessage.id)))) == 7
|
|
assert list(await session.scalars(select(OperationAudit.id))) == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_long_conversation_withdraw_loads_only_target_pair_and_latest_page(withdraw_db):
|
|
sessions, engine, start = withdraw_db
|
|
async with sessions() as session:
|
|
session.add_all(
|
|
[
|
|
ConversationMessage(
|
|
id=message_id,
|
|
conversation_id=1,
|
|
role="user",
|
|
content="较长的历史文字" * 100,
|
|
created_at=start + timedelta(minutes=message_id),
|
|
)
|
|
for message_id in range(8, 1008)
|
|
]
|
|
)
|
|
await session.commit()
|
|
async with sessions() as session:
|
|
user = await session.get(UserAccount, 1)
|
|
statements = []
|
|
|
|
def record_query(_, __, statement, ___, ____, _____):
|
|
statements.append(statement)
|
|
|
|
event.listen(engine.sync_engine, "before_cursor_execute", record_query)
|
|
try:
|
|
result = await ConversationService(session).withdraw_turn(1, 1, user)
|
|
finally:
|
|
event.remove(engine.sync_engine, "before_cursor_execute", record_query)
|
|
assert result.messages_total == 1004
|
|
assert [message.id for message in result.messages] == list(range(988, 1008))
|
|
assert result.has_more_messages is True
|
|
conversation = await session.get(Conversation, 1)
|
|
assert "messages" in inspect(conversation).unloaded
|
|
assert sum(isinstance(row, ConversationMessage) for row in session.identity_map.values()) <= 22
|
|
# Ownership + target, next message, count, latest page: no redundant header query.
|
|
assert len([sql for sql in statements if sql.lstrip().upper().startswith("SELECT")]) == 5
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_withdraw_rollback_restores_messages_and_drops_audit(withdraw_db):
|
|
sessions, _, _ = withdraw_db
|
|
async with sessions() as session:
|
|
await ConversationService(session).withdraw_turn(1, 1, await session.get(UserAccount, 1))
|
|
await session.rollback()
|
|
async with sessions() as session:
|
|
assert len(list(await session.scalars(select(ConversationMessage.id)))) == 7
|
|
assert list(await session.scalars(select(OperationAudit.id))) == []
|