tests/test_pr_area.py

import json

from app.config import DATA_DIR
from app.tasks.pr_area import AREAS, TASK, parse_output, score


def test_dataset_loaded_and_labels_valid():
    records = [json.loads(line) for line in (DATA_DIR / "prs.jsonl").open()]
    assert len(records) > 800
    assert {r["area"] for r in records} <= set(AREAS)
    numbers = [r["number"] for r in records]
    assert len(numbers) == len(set(numbers))


def test_sample_and_truth_roundtrip():
    s = TASK.sample()
    truth = TASK.lookup_truth(s.input_id)
    assert truth and truth["area"] in AREAS
    assert s.text.startswith("Title: ")
    assert TASK.lookup_truth("pr-999999999") is None


def test_parse_output_variants():
    assert parse_output('{"area": "docs"}')["area"] == "docs"
    assert parse_output('Sure!\n{"area": "Infra"}\nDone.')["area"] == "infra"
    assert parse_output('{"area": "not-a-label"}')["area"] is None
    assert parse_output("no json here")["area"] is None
    assert parse_output('{"wrong_key": "docs"}')["area"] is None


def test_score():
    truth = {"area": "docs", "pr_number": 1}
    assert score(truth, parse_output('{"area": "docs"}')) == {
        "format_ok": 1.0, "exact_match": 1.0,
    }
    assert score(truth, parse_output('{"area": "infra"}')) == {
        "format_ok": 1.0, "exact_match": 0.0,
    }
    assert score(truth, parse_output("garbage")) == {
        "format_ok": 0.0, "exact_match": 0.0,
    }