File size: 856 Bytes
88e15cd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 | 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"]
|