AgentEvalTool/tests/unit/test_intelligent_eval_task_queue.py

"""Unit tests for intelligent eval task queue."""

from datetime import timedelta

import pytest
from sqlmodel import Session, select

from agenteval.intelligent_eval import task_queue
from agenteval.intelligent_eval.models import IntelligentEvalStatus
from agenteval.storage.db import (
    IntelligentEvalDB,
    IntelligentEvalSessionDB,
    IntelligentEvalTaskQueueDB,
    utc_now,
)


def test_is_slot_due():
    """Test time slot due detection."""
    # Slot "8-10h" should be due after 8 hours
    slot = {"time_slot": "8-10h", "sessions": 2}
    assert task_queue._is_slot_due(slot, timedelta(hours=7)) is False
    assert task_queue._is_slot_due(slot, timedelta(hours=8)) is True
    assert task_queue._is_slot_due(slot, timedelta(hours=9)) is True

    # Invalid slot format
    assert task_queue._is_slot_due({"time_slot": "invalid"}, timedelta(hours=1)) is False
    assert task_queue._is_slot_due({}, timedelta(hours=1)) is False


def test_calculate_session_deficit(db_session: Session):
    """Test session deficit calculation."""
    # Create eval with plan
    eval_db = IntelligentEvalDB(
        name="test",
        target_id="target1",
        status=IntelligentEvalStatus.EXECUTING.value,
        started_at=utc_now() - timedelta(hours=9),
    )
    eval_db.set_plan({
        "time_distribution": [
            {"time_slot": "0-2h", "sessions": 1},
            {"time_slot": "8-10h", "sessions": 2},
        ],
        "estimated_sessions": 3,
    })
    db_session.add(eval_db)
    db_session.commit()

    # No sessions yet, should have 3 (1 from 0-2h, 2 from 8-10h)
    deficit = task_queue._calculate_session_deficit(eval_db, db_session)
    assert deficit == 3

    # Add 1 session
    session_db = IntelligentEvalSessionDB(
        eval_id=eval_db.id,
        target_id="target1",
        status="completed",
    )
    db_session.add(session_db)
    db_session.commit()

    # Should have 3, has 1, deficit = 2
    deficit = task_queue._calculate_session_deficit(eval_db, db_session)
    assert deficit == 2


def test_calculate_priority(db_session: Session):
    """Test task priority calculation."""
    # Eval with due slot and deficit
    eval_db = IntelligentEvalDB(
        name="test",
        target_id="target1",
        status=IntelligentEvalStatus.EXECUTING.value,
        started_at=utc_now() - timedelta(hours=9),
    )
    eval_db.set_plan({
        "time_distribution": [{"time_slot": "8-10h", "sessions": 2}],
        "estimated_sessions": 2,
    })
    db_session.add(eval_db)
    db_session.commit()

    priority = task_queue._calculate_priority(eval_db, db_session)
    # Base 100 - 50 (slot due) - 20 (deficit 2 * 10) - 20 (wait 9h / 10min = 54, capped at 20)
    assert priority == 10


def test_get_attention_reason(db_session: Session):
    """Test attention reason detection."""
    # Eval with due slot
    eval_db = IntelligentEvalDB(
        name="test",
        target_id="target1",
        status=IntelligentEvalStatus.EXECUTING.value,
        started_at=utc_now() - timedelta(hours=9),
    )
    eval_db.set_plan({
        "time_distribution": [{"time_slot": "8-10h", "sessions": 2}],
        "estimated_sessions": 2,
    })
    db_session.add(eval_db)
    db_session.commit()

    reason = task_queue._get_attention_reason(eval_db, db_session)
    assert reason == "slot_due"

    # Add all sessions as completed
    for _ in range(2):
        session_db = IntelligentEvalSessionDB(
            eval_id=eval_db.id,
            target_id="target1",
            status="completed",
        )
        db_session.add(session_db)
    db_session.commit()

    reason = task_queue._get_attention_reason(eval_db, db_session)
    assert reason == "all_sessions_completed"


def test_has_pending_task(db_session: Session):
    """Test pending task detection (去重)."""
    eval_db = IntelligentEvalDB(
        name="test",
        target_id="target1",
        status=IntelligentEvalStatus.EXECUTING.value,
    )
    db_session.add(eval_db)
    db_session.commit()

    # No pending task initially
    assert task_queue._has_pending_task(eval_db.id, db_session) is False

    # Add pending task
    task = IntelligentEvalTaskQueueDB(
        eval_id=eval_db.id,
        status="pending",
        priority=1,
        reason="slot_due",
    )
    db_session.add(task)
    db_session.commit()

    # Should detect pending task
    assert task_queue._has_pending_task(eval_db.id, db_session) is True


def test_scan_and_enqueue_tasks(db_session: Session):
    """Test task scanning and enqueueing."""
    # Create eval that needs attention
    eval_db = IntelligentEvalDB(
        name="test",
        target_id="target1",
        status=IntelligentEvalStatus.EXECUTING.value,
        started_at=utc_now() - timedelta(hours=9),
    )
    eval_db.set_plan({
        "time_distribution": [{"time_slot": "8-10h", "sessions": 2}],
        "estimated_sessions": 2,
    })
    db_session.add(eval_db)
    db_session.commit()

    # Scan and enqueue
    enqueued = task_queue.scan_and_enqueue_tasks(db_session)
    assert enqueued == 1

    # Verify task created
    task = db_session.exec(
        select(IntelligentEvalTaskQueueDB).where(IntelligentEvalTaskQueueDB.eval_id == eval_db.id)
    ).first()
    assert task is not None
    assert task.status == "pending"
    assert task.reason == "slot_due"
    assert task.priority < 100  # Should have reduced priority

    # Scan again, should not create duplicate
    enqueued = task_queue.scan_and_enqueue_tasks(db_session)
    assert enqueued == 0


def test_get_next_task(db_session: Session):
    """Test getting next task (highest priority)."""
    # Create tasks with different priorities
    task1 = IntelligentEvalTaskQueueDB(
        eval_id="eval1",
        status="pending",
        priority=50,
        reason="slot_due",
    )
    task2 = IntelligentEvalTaskQueueDB(
        eval_id="eval2",
        status="pending",
        priority=10,  # Higher priority (lower number)
        reason="slot_due",
    )
    task3 = IntelligentEvalTaskQueueDB(
        eval_id="eval3",
        status="assigned",  # Not pending
        priority=1,
        reason="slot_due",
    )
    db_session.add_all([task1, task2, task3])
    db_session.commit()

    # Should get task2 (priority 10)
    next_task = task_queue.get_next_task(db_session)
    assert next_task is not None
    assert next_task.id == task2.id


def test_assign_task(db_session: Session):
    """Test task assignment."""
    task = IntelligentEvalTaskQueueDB(
        eval_id="eval1",
        status="pending",
        priority=1,
        reason="slot_due",
    )
    db_session.add(task)
    db_session.commit()

    # Assign task
    success = task_queue.assign_task(task.id, "cron1", db_session)
    assert success is True

    # Verify assignment
    db_session.refresh(task)
    assert task.status == "assigned"
    assert task.assigned_cron_id == "cron1"
    assert task.assigned_at is not None

    # Try to assign again (should fail)
    success = task_queue.assign_task(task.id, "cron2", db_session)
    assert success is False


def test_complete_task(db_session: Session):
    """Test task completion."""
    task = IntelligentEvalTaskQueueDB(
        eval_id="eval1",
        status="assigned",
        priority=1,
        reason="slot_due",
        assigned_cron_id="cron1",
    )
    db_session.add(task)
    db_session.commit()

    # Complete task successfully
    success = task_queue.complete_task(task.id, True, None, db_session)
    assert success is True

    # Verify completion
    db_session.refresh(task)
    assert task.status == "completed"
    assert task.completed_at is not None
    assert task.error is None

    # Complete task with error
    task2 = IntelligentEvalTaskQueueDB(
        eval_id="eval2",
        status="assigned",
        priority=1,
        reason="slot_due",
        assigned_cron_id="cron1",
    )
    db_session.add(task2)
    db_session.commit()

    success = task_queue.complete_task(task2.id, False, "test error", db_session)
    assert success is True

    db_session.refresh(task2)
    assert task2.status == "failed"
    assert task2.error == "test error"