AgentEvalTool/tests/integration/test_intelligent_eval_task_queue_api.py

"""Integration tests for intelligent eval task queue API."""

from datetime import timedelta

import pytest
from fastapi.testclient import TestClient
from sqlmodel import Session, SQLModel, create_engine, select

from agenteval.intelligent_eval.models import IntelligentEvalStatus
from agenteval.storage.db import IntelligentEvalDB, IntelligentEvalTaskQueueDB, utc_now
from agenteval.web.app import app
from agenteval.web.deps import get_db


@pytest.fixture()
def client(tmp_path):
    """Create a TestClient with a fresh database."""
    from agenteval.storage.db import (  # noqa: F401
        IntelligentEvalDB,
        IntelligentEvalSessionDB,
        IntelligentEvalTaskQueueDB,
    )

    engine = create_engine(
        f"sqlite:///{tmp_path / 'test.db'}",
        connect_args={"check_same_thread": False},
    )
    SQLModel.metadata.create_all(engine)
    session = Session(engine)

    def override_get_db():
        try:
            yield session
        finally:
            pass

    app.dependency_overrides[get_db] = override_get_db
    client = TestClient(app)
    yield client
    app.dependency_overrides.clear()
    session.close()
    engine.dispose()


@pytest.fixture()
def db_session(client):
    """Get the database session from the client fixture."""
    # The session is stored in the dependency override
    return next(app.dependency_overrides[get_db]())


def test_get_next_task_empty(client: TestClient):
    """Test getting next task when queue is empty."""
    response = client.get("/api/intelligent-evals/tasks/next")
    assert response.status_code == 200
    assert response.json() == {"task": None}


def test_get_next_task_with_pending_task(client: TestClient, db_session: Session):
    """Test getting next task when there is a pending task."""
    # Create eval
    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()

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

    # Get next task
    response = client.get("/api/intelligent-evals/tasks/next")
    assert response.status_code == 200

    data = response.json()
    assert data["task"] is not None
    assert data["task"]["id"] == task.id
    assert data["task"]["eval_id"] == eval_db.id
    assert data["task"]["priority"] == 10
    assert data["task"]["reason"] == "slot_due"
    assert data["task"]["eval"]["id"] == eval_db.id
    assert data["task"]["eval"]["name"] == "test"
    assert data["task"]["eval"]["status"] == IntelligentEvalStatus.EXECUTING.value


def test_assign_task(client: TestClient, db_session: Session):
    """Test assigning a task to a cron."""
    # Create pending task
    task = IntelligentEvalTaskQueueDB(
        eval_id="eval1",
        status="pending",
        priority=1,
        reason="slot_due",
    )
    db_session.add(task)
    db_session.commit()

    # Assign task
    response = client.post(f"/api/intelligent-evals/tasks/{task.id}/assign?cron_id=cron1")
    assert response.status_code == 200
    assert response.json() == {"success": True}

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


def test_assign_task_not_found(client: TestClient):
    """Test assigning a non-existent task."""
    response = client.post("/api/intelligent-evals/tasks/nonexistent/assign?cron_id=cron1")
    assert response.status_code == 404


def test_complete_task(client: TestClient, db_session: Session):
    """Test completing a task."""
    # Create assigned task
    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
    response = client.post(f"/api/intelligent-evals/tasks/{task.id}/complete?success=true")
    assert response.status_code == 200
    assert response.json() == {"success": True}

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


def test_complete_task_with_error(client: TestClient, db_session: Session):
    """Test completing a task with error."""
    # Create assigned task
    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 with error
    response = client.post(f"/api/intelligent-evals/tasks/{task.id}/complete?success=false&error=test_error")
    assert response.status_code == 200
    assert response.json() == {"success": True}

    # Verify completion
    db_session.refresh(task)
    assert task.status == "failed"
    assert task.error == "test_error"


def test_end_to_end_task_lifecycle(client: TestClient, db_session: Session):
    """Test end-to-end task lifecycle: create eval -> scan -> enqueue -> assign -> complete."""
    # 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 tasks
    from agenteval.intelligent_eval import task_queue

    enqueued = task_queue.scan_and_enqueue_tasks(db_session)
    assert enqueued == 1

    # Get next task
    response = client.get("/api/intelligent-evals/tasks/next")
    assert response.status_code == 200
    task_data = response.json()["task"]
    assert task_data is not None
    assert task_data["eval_id"] == eval_db.id

    # Assign task
    response = client.post(f"/api/intelligent-evals/tasks/{task_data['id']}/assign?cron_id=cron1")
    assert response.status_code == 200

    # Complete task
    response = client.post(f"/api/intelligent-evals/tasks/{task_data['id']}/complete?success=true")
    assert response.status_code == 200

    # Verify task completed
    task = db_session.get(IntelligentEvalTaskQueueDB, task_data["id"])
    assert task.status == "completed"