from pathlib import Path

import pytest

from fugu_lite.config import load_worker_pool
from fugu_lite.generate import generate_reward_dataset
from fugu_lite.io import load_rewards, load_tasks


@pytest.mark.asyncio
async def test_mock_reward_generation(tmp_path: Path):
    root = Path(__file__).parents[1]
    tasks = load_tasks(root / "data/tasks.sample.jsonl")
    selected = [tasks[0], tasks[6], tasks[12]]
    config = load_worker_pool(root / "configs/workers.mock.yaml")
    output = tmp_path / "rewards.jsonl"
    await generate_reward_dataset(selected, config, output)
    records = load_rewards(output)
    assert len(records) == 3
    winners = [
        record.worker_ids[max(range(len(record.rewards)), key=record.rewards.__getitem__)]
        for record in records
    ]
    assert winners == ["math_expert", "code_expert", "general_expert"]