0cca4963d1
Both the single-run path and the campaign scheduler drove long-lived asyncio tasks through their own duplicated _tasks/_cancel_events dicts and shutdown loops. Collapse them into one deep TaskRegistry module, instantiated as run_registry and campaign_registry. launch() creates the cancel event before the task (so a cancel during startup is never lost), wires done-callback cleanup, and is idempotent per id; this makes runs.py's hard-cancel fallback provably dead, so it is removed. App shutdown now gracefully stops in-flight runs too, not just campaigns.
155 lines
6.0 KiB
Python
155 lines
6.0 KiB
Python
"""Integration tests for the durable campaign scheduler loop (ticket 03).
|
||
|
||
Drives ``campaign_runner.start_campaign`` 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_runner import request_cancel, start_campaign
|
||
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))
|
||
|
||
|
||
async def _await_task(campaign_id, timeout=3.0):
|
||
import asyncio
|
||
task = campaign_runner.campaign_registry.get(campaign_id)
|
||
if task is not None:
|
||
await asyncio.wait_for(task, timeout=timeout)
|
||
|
||
|
||
async def test_loop_runs_to_completion(seeded_db):
|
||
campaign = _make_campaign(seeded_db)
|
||
start_campaign(campaign.id, seeded_db, tick_seconds=TICK)
|
||
await _await_task(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_restart_recovery_does_not_respawn(seeded_db):
|
||
# 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.
|
||
start_campaign(campaign.id, seeded_db, tick_seconds=TICK)
|
||
await _await_task(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):
|
||
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).
|
||
start_campaign(campaign.id, seeded_db, tick_seconds=TICK)
|
||
request_cancel(campaign.id) # stop the long-window loop promptly
|
||
await _await_task(campaign.id)
|
||
|
||
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):
|
||
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),
|
||
],
|
||
)
|
||
start_campaign(campaign.id, seeded_db, tick_seconds=TICK)
|
||
await asyncio.sleep(0.05) # let the first tick spawn entry 0
|
||
|
||
# Cancel: mark DB authoritative + signal the live loop.
|
||
repo = CampaignRepository(seeded_db)
|
||
current = repo.get(campaign.id)
|
||
current.status = CampaignStatus.CANCELLED
|
||
current.completed_at = utc_now()
|
||
repo.update(current)
|
||
request_cancel(campaign.id)
|
||
await _await_task(campaign.id)
|
||
|
||
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
|