61 lines
2.4 KiB
Python
61 lines
2.4 KiB
Python
from datetime import UTC, datetime, timedelta
|
|
from uuid import uuid4
|
|
|
|
import pytest
|
|
from sqlalchemy import delete, update
|
|
|
|
from app.infrastructure.db import SessionFactory
|
|
from app.model.conversation import ConversationMessage
|
|
from app.model.platform import AgentRun, RequestIdempotency
|
|
from app.repository.agent_run_repository import AgentRunRepository
|
|
|
|
|
|
@pytest.mark.integration
|
|
@pytest.mark.asyncio
|
|
async def test_worker_lease_renew_and_takeover() -> None:
|
|
now = datetime.now(UTC).replace(tzinfo=None)
|
|
session_id = f"lease-{uuid4()}"
|
|
run_id = str(uuid4())
|
|
trace_id = str(uuid4())
|
|
async with SessionFactory() as session:
|
|
idem = RequestIdempotency(
|
|
user_id=1, session_id=session_id, agent_type="customer_service",
|
|
idempotency_key=f"lease-key-{uuid4()}", request_hash="a" * 64,
|
|
trace_id=trace_id, status="processing", expire_at=now + timedelta(hours=1),
|
|
created_at=now, updated_at=now,
|
|
)
|
|
message = ConversationMessage(
|
|
session_id=session_id, customer_id=1, portal="api", role="user",
|
|
content="lease test", trace_id=trace_id, created_at=now,
|
|
)
|
|
session.add_all([idem, message])
|
|
await session.flush()
|
|
run = AgentRun(
|
|
run_id=run_id, idempotency_id=idem.id, session_id=session_id, user_id=1,
|
|
agent_type="customer_service", trace_id=trace_id, request_message_id=message.id,
|
|
created_at=now, updated_at=now,
|
|
)
|
|
session.add(run)
|
|
await session.commit()
|
|
try:
|
|
repo = AgentRunRepository(session)
|
|
assert await repo.claim(run_id, "worker-a", 60)
|
|
assert await repo.renew(run_id, "worker-a", 60)
|
|
assert not await repo.claim(run_id, "worker-b", 60)
|
|
await session.execute(
|
|
update(AgentRun).where(AgentRun.run_id == run_id).values(
|
|
locked_until=now - timedelta(seconds=1)
|
|
)
|
|
)
|
|
await session.commit()
|
|
assert await repo.claim(run_id, "worker-b", 60)
|
|
finally:
|
|
await session.execute(delete(AgentRun).where(AgentRun.run_id == run_id))
|
|
await session.execute(
|
|
delete(RequestIdempotency).where(RequestIdempotency.id == idem.id)
|
|
)
|
|
await session.execute(
|
|
delete(ConversationMessage).where(ConversationMessage.session_id == session_id)
|
|
)
|
|
await session.commit()
|