AgentEvalTool/tests/integration/test_campaign_scheduler_loop.py

"""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