78 lines
3.0 KiB
Python
78 lines
3.0 KiB
Python
"""Tests for atomic Campaign cancellation and settlement ordering."""
|
|
|
|
import pytest
|
|
from agenteval.evaluation.campaign_lifecycle import CampaignLifecycleError, cancel_campaign
|
|
from agenteval.exploration.models import ExplorationSession, ExplorationSessionStatus
|
|
from agenteval.models import Campaign, CampaignPlanEntry, CampaignStatus
|
|
from agenteval.storage.repository import CampaignRepository, ExplorationSessionRepository
|
|
from sqlalchemy import event
|
|
|
|
|
|
def _campaign(session, status=CampaignStatus.RUNNING):
|
|
return CampaignRepository(session).create(
|
|
Campaign(
|
|
id="campaign-1",
|
|
name="campaign",
|
|
target_id="t-1",
|
|
status=status,
|
|
window_seconds=60,
|
|
plan=[CampaignPlanEntry(scenario_id="s-1", offset_seconds=0, count=1)],
|
|
)
|
|
)
|
|
|
|
|
|
def test_cancel_is_atomic_and_stop_happens_after_durable_state(db_session):
|
|
_campaign(db_session)
|
|
observed = []
|
|
|
|
def stop(campaign_id):
|
|
observed.append(CampaignRepository(db_session).get(campaign_id).status)
|
|
|
|
cancelled = cancel_campaign(db_session, "campaign-1", stop=stop)
|
|
|
|
assert cancelled.status is CampaignStatus.CANCELLED
|
|
assert cancelled.completed_at is not None
|
|
assert observed == [CampaignStatus.CANCELLED]
|
|
|
|
|
|
def test_cancel_distinguishes_missing_and_terminal_campaign(db_session):
|
|
with pytest.raises(CampaignLifecycleError) as missing:
|
|
cancel_campaign(db_session, "missing")
|
|
assert missing.value.status_code == 404
|
|
|
|
_campaign(db_session, CampaignStatus.COMPLETED)
|
|
with pytest.raises(CampaignLifecycleError) as terminal:
|
|
cancel_campaign(db_session, "campaign-1")
|
|
assert terminal.value.status_code == 400
|
|
|
|
|
|
def test_cancel_is_idempotency_guarded_by_status(db_session):
|
|
_campaign(db_session)
|
|
cancel_campaign(db_session, "campaign-1")
|
|
with pytest.raises(CampaignLifecycleError) as again:
|
|
cancel_campaign(db_session, "campaign-1")
|
|
assert again.value.status_code == 400
|
|
|
|
|
|
def test_cancel_settlement_failure_rolls_back_campaign_and_session(db_session):
|
|
_campaign(db_session)
|
|
session_obj = ExplorationSessionRepository(db_session).create(
|
|
ExplorationSession(campaign_id="campaign-1", target_id="t-1")
|
|
)
|
|
engine = db_session.get_bind()
|
|
|
|
def fail_exploration_update(_connection, _cursor, statement, _parameters, _context, _executemany):
|
|
if statement.lstrip().upper().startswith("UPDATE EXPLORATION_SESSIONS"):
|
|
raise RuntimeError("settlement failed")
|
|
|
|
event.listen(engine, "before_cursor_execute", fail_exploration_update)
|
|
try:
|
|
with pytest.raises(RuntimeError, match="settlement failed"):
|
|
cancel_campaign(db_session, "campaign-1")
|
|
finally:
|
|
event.remove(engine, "before_cursor_execute", fail_exploration_update)
|
|
db_session.rollback()
|
|
|
|
assert CampaignRepository(db_session).get("campaign-1").status is CampaignStatus.RUNNING
|
|
assert ExplorationSessionRepository(db_session).get(session_obj.id).status is ExplorationSessionStatus.RUNNING
|