234 lines
8.1 KiB
Python
234 lines
8.1 KiB
Python
"""Integration tests for the durable campaign scheduler loop (ticket 03).
|
||
|
||
Drives ``CampaignRuntime`` with a tiny tick and a compressed
|
||
time scale so a whole window elapses in a few real milliseconds. Covers:
|
||
auto-run-to-completion, restart recovery (no double-spawn, no lost progress),
|
||
and cancellation (no further spawning).
|
||
"""
|
||
|
||
from datetime import timedelta
|
||
|
||
import pytest
|
||
from agenteval.evaluation import campaign_runner
|
||
from agenteval.evaluation.campaign_lifecycle import cancel_campaign
|
||
from agenteval.evaluation.campaign_runner import CampaignRuntime
|
||
from agenteval.models import (
|
||
Campaign,
|
||
CampaignPlanEntry,
|
||
CampaignStatus,
|
||
Case,
|
||
CaseType,
|
||
ChannelType,
|
||
EvalTarget,
|
||
PlatformType,
|
||
RunStatus,
|
||
Scenario,
|
||
TargetStatus,
|
||
)
|
||
from agenteval.storage.db import utc_now
|
||
from agenteval.storage.repository import CampaignRepository, RunRepository, ScenarioRepository, TargetRepository
|
||
|
||
from tests.unit.mock_channel import MockChannel
|
||
|
||
TICK = 0.01
|
||
|
||
|
||
@pytest.fixture()
|
||
def seeded_db(db_session, monkeypatch):
|
||
from agenteval.channels import factory as factory_module
|
||
from agenteval.evaluation import engine as engine_module
|
||
from agenteval.storage import db as db_module
|
||
from agenteval.storage import repository as repo_module
|
||
|
||
def _test_get_session():
|
||
return db_session
|
||
|
||
monkeypatch.setattr(db_module, "get_session", _test_get_session)
|
||
monkeypatch.setattr(repo_module, "get_session", _test_get_session)
|
||
monkeypatch.setattr(engine_module, "get_session", _test_get_session)
|
||
monkeypatch.setattr(campaign_runner, "get_session", _test_get_session)
|
||
|
||
channel = MockChannel(reply_delay=0.0)
|
||
monkeypatch.setattr(factory_module.ChannelFactory, "create", lambda target: channel)
|
||
|
||
TargetRepository(db_session).create(
|
||
EvalTarget(
|
||
id="t-1",
|
||
name="mock-target",
|
||
platform=PlatformType.AI_DIGITAL_EMPLOYEE,
|
||
channel_type=ChannelType.TUTU_API,
|
||
channel_config={"base_url": "http://mock", "token": "x"},
|
||
status=TargetStatus.ACTIVE,
|
||
)
|
||
)
|
||
ScenarioRepository(db_session).create(
|
||
Scenario(
|
||
id="s-1",
|
||
name="mock-scenario",
|
||
cases=[Case(id="c1", type=CaseType.SINGLE, messages=["hi"])],
|
||
)
|
||
)
|
||
return db_session
|
||
|
||
|
||
def _make_campaign(session, **overrides) -> Campaign:
|
||
payload = dict(
|
||
name="loop",
|
||
target_id="t-1",
|
||
window_seconds=1,
|
||
time_scale=1000.0, # elapsed×1000 → window ends in ~1ms of real time
|
||
plan=[
|
||
CampaignPlanEntry(scenario_id="s-1", offset_seconds=0, count=1),
|
||
CampaignPlanEntry(scenario_id="s-1", offset_seconds=1, count=1),
|
||
],
|
||
)
|
||
payload.update(overrides)
|
||
return CampaignRepository(session).create(Campaign(**payload))
|
||
|
||
|
||
@pytest.fixture()
|
||
async def runtime(seeded_db):
|
||
value = CampaignRuntime(session_factory=lambda: seeded_db, tick_seconds=TICK)
|
||
yield value
|
||
await value.shutdown()
|
||
|
||
|
||
async def _await_terminal(session, campaign_id, timeout=3.0):
|
||
import asyncio
|
||
|
||
async with asyncio.timeout(timeout):
|
||
while CampaignRepository(session).get(campaign_id).status is CampaignStatus.RUNNING:
|
||
await asyncio.sleep(TICK)
|
||
|
||
|
||
async def test_loop_runs_to_completion(seeded_db, runtime):
|
||
campaign = _make_campaign(seeded_db)
|
||
assert runtime.start(campaign.id)
|
||
await _await_terminal(seeded_db, campaign.id)
|
||
|
||
final = CampaignRepository(seeded_db).get(campaign.id)
|
||
assert final.status == CampaignStatus.COMPLETED
|
||
assert final.completed_at is not None
|
||
|
||
runs = RunRepository(seeded_db).list_by_campaign(campaign.id)
|
||
assert len(runs) == 2 # both plan entries spawned exactly once
|
||
assert all(r.status == RunStatus.COMPLETED for r in runs)
|
||
|
||
|
||
async def test_runtime_start_is_idempotent(seeded_db, runtime):
|
||
campaign = _make_campaign(seeded_db)
|
||
|
||
assert runtime.start(campaign.id)
|
||
assert runtime.start(campaign.id)
|
||
await _await_terminal(seeded_db, campaign.id)
|
||
|
||
assert len(RunRepository(seeded_db).list_by_campaign(campaign.id)) == 2
|
||
|
||
|
||
async def test_runtime_accepts_internal_clock_and_executor_adapters(seeded_db):
|
||
started_at = utc_now()
|
||
campaign = _make_campaign(seeded_db, started_at=started_at, status=CampaignStatus.RUNNING)
|
||
executed: list[tuple[int, int]] = []
|
||
|
||
async def execute_child(_campaign, _scenario, run, session):
|
||
executed.append((run.campaign_plan_index, run.campaign_occurrence_index))
|
||
run.status = RunStatus.COMPLETED
|
||
RunRepository(session).update(run)
|
||
|
||
runtime = CampaignRuntime(
|
||
session_factory=lambda: seeded_db,
|
||
now=lambda: started_at + timedelta(seconds=1),
|
||
tick_seconds=TICK,
|
||
execute_child_run=execute_child,
|
||
)
|
||
try:
|
||
assert runtime.start(campaign.id)
|
||
await _await_terminal(seeded_db, campaign.id)
|
||
finally:
|
||
await runtime.shutdown()
|
||
|
||
assert executed == [(0, 0), (1, 0)]
|
||
|
||
|
||
async def test_loop_retries_completion_after_settlement_failure(seeded_db, runtime, monkeypatch):
|
||
campaign = _make_campaign(seeded_db)
|
||
real_complete = campaign_runner.complete_campaign
|
||
attempts = 0
|
||
|
||
def flaky_complete(*args, **kwargs):
|
||
nonlocal attempts
|
||
attempts += 1
|
||
if attempts == 1:
|
||
raise RuntimeError("settlement failed")
|
||
return real_complete(*args, **kwargs)
|
||
|
||
monkeypatch.setattr(campaign_runner, "complete_campaign", flaky_complete)
|
||
assert runtime.start(campaign.id)
|
||
await _await_terminal(seeded_db, campaign.id)
|
||
|
||
assert attempts == 2
|
||
assert CampaignRepository(seeded_db).get(campaign.id).status is CampaignStatus.COMPLETED
|
||
|
||
|
||
async def test_restart_recovery_does_not_respawn(seeded_db, runtime):
|
||
# Simulate a campaign that was already RUNNING before a restart, with its
|
||
# window start well in the past and entry 0 already recorded as spawned.
|
||
campaign = _make_campaign(seeded_db)
|
||
campaign.status = CampaignStatus.RUNNING
|
||
campaign.started_at = utc_now() - timedelta(seconds=10)
|
||
campaign.summary = {"scheduler": {"spawned_indices": [0]}}
|
||
CampaignRepository(seeded_db).update(campaign)
|
||
|
||
# Recovery relaunches the loop for the already-RUNNING campaign.
|
||
recovery = runtime.recover()
|
||
assert recovery.resumed_campaigns == 1
|
||
await _await_terminal(seeded_db, campaign.id)
|
||
|
||
final = CampaignRepository(seeded_db).get(campaign.id)
|
||
assert final.status == CampaignStatus.COMPLETED
|
||
|
||
runs = RunRepository(seeded_db).list_by_campaign(campaign.id)
|
||
# Entry 0 was already spawned pre-restart (not re-spawned); only entry 1 runs.
|
||
assert len(runs) == 1
|
||
|
||
|
||
async def test_restart_preserves_original_window_start(seeded_db, runtime):
|
||
original = utc_now() - timedelta(seconds=5)
|
||
campaign = _make_campaign(seeded_db, window_seconds=100000, time_scale=1.0)
|
||
campaign.status = CampaignStatus.RUNNING
|
||
campaign.started_at = original
|
||
CampaignRepository(seeded_db).update(campaign)
|
||
|
||
# Resume must NOT reset started_at (that would rewind the window clock).
|
||
assert runtime.start(campaign.id)
|
||
cancel_campaign(seeded_db, campaign.id, stop=runtime.cancel)
|
||
|
||
reloaded = CampaignRepository(seeded_db).get(campaign.id)
|
||
assert abs((reloaded.started_at.replace(tzinfo=None) - original.replace(tzinfo=None)).total_seconds()) < 1
|
||
|
||
|
||
async def test_cancel_stops_further_spawning(seeded_db, runtime):
|
||
import asyncio
|
||
|
||
# Long real window (scale 1.0) so entry 1 (offset 50s) never comes due fast.
|
||
campaign = _make_campaign(
|
||
seeded_db,
|
||
window_seconds=100,
|
||
time_scale=1.0,
|
||
plan=[
|
||
CampaignPlanEntry(scenario_id="s-1", offset_seconds=0, count=1),
|
||
CampaignPlanEntry(scenario_id="s-1", offset_seconds=50, count=1),
|
||
],
|
||
)
|
||
assert runtime.start(campaign.id)
|
||
await asyncio.sleep(0.05) # let the first tick spawn entry 0
|
||
|
||
# Cancel: mark DB authoritative + signal the live loop.
|
||
repo = CampaignRepository(seeded_db)
|
||
cancel_campaign(seeded_db, campaign.id, stop=runtime.cancel)
|
||
|
||
final = repo.get(campaign.id)
|
||
assert final.status == CampaignStatus.CANCELLED
|
||
runs = RunRepository(seeded_db).list_by_campaign(campaign.id)
|
||
assert len(runs) == 1 # entry 1 (offset 50s) never spawned
|