38 lines
1.3 KiB
Python
38 lines
1.3 KiB
Python
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from scripts.generate_codegen import build_prompt, generate_predictions, read_completed_task_ids
|
|
|
|
|
|
def test_generate_predictions_groups_samples_by_task():
|
|
task = SimpleNamespace(task_id="add", entry_point="add", prompt="Implement add(a, b).")
|
|
prompts = []
|
|
|
|
def generate(prompt):
|
|
prompts.append(prompt)
|
|
return f"candidate-{len(prompts)}"
|
|
|
|
records = generate_predictions([task], generate, samples_per_task=3)
|
|
|
|
assert records == [{"task_id": "add", "candidates": ["candidate-1", "candidate-2", "candidate-3"]}]
|
|
assert all("Required entry point: add" in prompt for prompt in prompts)
|
|
assert "exactly one ```python``` code block" in build_prompt(task)
|
|
|
|
|
|
def test_generate_predictions_rejects_empty_sample_count():
|
|
with pytest.raises(ValueError, match="at least 1"):
|
|
generate_predictions([], lambda prompt: prompt, samples_per_task=0)
|
|
|
|
|
|
def test_read_completed_task_ids_supports_resume(tmp_path):
|
|
output = tmp_path / "predictions.jsonl"
|
|
output.write_text(
|
|
'{"task_id":"one","candidates":["a"]}\n'
|
|
'{"task_id":"two","candidates":["b"]}\n',
|
|
encoding="utf-8",
|
|
)
|
|
|
|
assert read_completed_task_ids(output) == {"one", "two"}
|
|
assert read_completed_task_ids(tmp_path / "missing.jsonl") == set()
|