Video-Tree-TRM5/tests/unit/test_question_loader.py

"""app/question_gen/loader.py 单元测试。"""

from __future__ import annotations

import json
from pathlib import Path

import pytest

from app.question_gen.loader import load_benchmark, stratified_sample
from core.types import GeneratedQuestion


@pytest.fixture()
def benchmark_dir(tmp_path: Path) -> Path:
    """创建包含 benchmark JSON 的临时目录。"""
    data = [
        {
            "question_id": "1-1",
            "task_type": "Action Reasoning",
            "question": "What happened?",
            "options": ["A. X", "B. Y", "C. Z", "D. W"],
            "answer": "A",
        },
        {
            "question_id": "1-2",
            "task_type": "OCR Problems",
            "question": "What text is shown?",
            "options": ["A. Hello", "B. World", "C. Foo", "D. Bar"],
            "answer": "B",
        },
    ]
    (tmp_path / "video_abc.json").write_text(json.dumps(data), encoding="utf-8")
    return tmp_path


class TestLoadBenchmark:
    def test_loads_questions_from_json(self, benchmark_dir: Path) -> None:
        questions = load_benchmark(benchmark_dir)
        assert len(questions) == 2

    def test_video_id_from_filename(self, benchmark_dir: Path) -> None:
        questions = load_benchmark(benchmark_dir)
        assert all(q.video_id == "video_abc" for q in questions)

    def test_fields_mapped_correctly(self, benchmark_dir: Path) -> None:
        questions = load_benchmark(benchmark_dir)
        q = questions[0]
        assert q.question_id == "1-1"
        assert q.task_type == "Action Reasoning"
        assert q.question == "What happened?"
        assert q.options == ("A. X", "B. Y", "C. Z", "D. W")
        assert q.answer == "A"

    def test_options_is_tuple(self, benchmark_dir: Path) -> None:
        questions = load_benchmark(benchmark_dir)
        assert isinstance(questions[0].options, tuple)

    def test_source_nodes_is_empty_tuple(self, benchmark_dir: Path) -> None:
        questions = load_benchmark(benchmark_dir)
        assert questions[0].source_nodes == ()

    def test_difficulty_defaults_to_medium_for_legacy(self, benchmark_dir: Path) -> None:
        questions = load_benchmark(benchmark_dir)
        assert questions[0].difficulty == "medium"

    def test_difficulty_from_json_when_present(self, tmp_path: Path) -> None:
        data = [
            {
                "question_id": "2-1",
                "task_type": "OCR Problems",
                "question": "Q?",
                "options": ["A. 1", "B. 2", "C. 3", "D. 4"],
                "answer": "C",
                "difficulty": "hard",
            }
        ]
        (tmp_path / "vid.json").write_text(json.dumps(data), encoding="utf-8")
        questions = load_benchmark(tmp_path)
        assert questions[0].difficulty == "hard"

    def test_empty_directory_returns_empty_list(self, tmp_path: Path) -> None:
        questions = load_benchmark(tmp_path)
        assert questions == []

    def test_sorted_by_filename(self, tmp_path: Path) -> None:
        for name in ["z_video.json", "a_video.json"]:
            data = [
                {
                    "question_id": f"{name}-1",
                    "task_type": "T",
                    "question": "Q?",
                    "options": ["A", "B", "C", "D"],
                    "answer": "A",
                }
            ]
            (tmp_path / name).write_text(json.dumps(data), encoding="utf-8")
        questions = load_benchmark(tmp_path)
        assert questions[0].video_id == "a_video"
        assert questions[1].video_id == "z_video"

    def test_returns_generated_question_instances(self, benchmark_dir: Path) -> None:
        questions = load_benchmark(benchmark_dir)
        assert all(isinstance(q, GeneratedQuestion) for q in questions)

    def test_loads_real_benchmark(self) -> None:
        """使用真实 benchmark 数据验证加载正确性。"""
        real_dir = Path("store/questions/benchmarks/Video-MME")
        if not real_dir.exists():
            pytest.skip("真实 benchmark 数据不存在")
        questions = load_benchmark(real_dir)
        assert len(questions) > 0
        for q in questions:
            assert isinstance(q, GeneratedQuestion)
            assert len(q.options) == 4
            assert q.answer in ("A", "B", "C", "D")

    def test_malformed_json_raises(self, tmp_path: Path) -> None:
        """非法 JSON 文件应抛出 json.JSONDecodeError。"""
        (tmp_path / "bad.json").write_text("not valid json{{{", encoding="utf-8")
        with pytest.raises(json.JSONDecodeError):
            load_benchmark(tmp_path)

    def test_missing_required_field_raises(self, tmp_path: Path) -> None:
        """缺少必需字段（如 question_id）应抛出 KeyError。"""
        data = [{"task_type": "T", "question": "Q?", "options": ["A"], "answer": "A"}]
        (tmp_path / "vid.json").write_text(json.dumps(data), encoding="utf-8")
        with pytest.raises(KeyError):
            load_benchmark(tmp_path)


def _make_questions(n: int, task_type: str = "T") -> list[GeneratedQuestion]:
    """辅助函数：批量构造题目。"""
    return [
        GeneratedQuestion(
            question_id=f"{task_type}-{i}",
            video_id="v1",
            task_type=task_type,
            question=f"Q{i}?",
            options=("A", "B", "C", "D"),
            answer="A",
            source_nodes=(),
            difficulty="medium",
        )
        for i in range(n)
    ]


class TestStratifiedSample:
    def test_natural_distribution(self) -> None:
        questions = _make_questions(20)
        result = stratified_sample(
            questions=questions,
            correctness={},
            size=10,
            correct_ratio=None,
            task_types=None,
            seed=42,
            min_per_class=None,
        )
        assert len(result) == 10

    def test_natural_distribution_pool_insufficient(self) -> None:
        questions = _make_questions(5)
        with pytest.raises(ValueError, match="自然分布采样不足"):
            stratified_sample(
                questions=questions,
                correctness={},
                size=10,
                correct_ratio=None,
                task_types=None,
                seed=42,
                min_per_class=None,
            )

    def test_ratio_stratified(self) -> None:
        questions = _make_questions(20)
        correctness = {f"T-{i}": i < 10 for i in range(20)}
        result = stratified_sample(
            questions=questions,
            correctness=correctness,
            size=10,
            correct_ratio=0.6,
            task_types=None,
            seed=42,
            min_per_class=None,
        )
        assert len(result) == 10
        correct_count = sum(1 for q in result if correctness.get(q.question_id, False))
        assert correct_count == 6

    def test_ratio_stratified_correct_first(self) -> None:
        questions = _make_questions(20)
        correctness = {f"T-{i}": i < 10 for i in range(20)}
        result = stratified_sample(
            questions=questions,
            correctness=correctness,
            size=10,
            correct_ratio=0.5,
            task_types=None,
            seed=42,
            min_per_class=None,
        )
        n_correct = round(10 * 0.5)
        for q in result[:n_correct]:
            assert correctness.get(q.question_id, False) is True
        for q in result[n_correct:]:
            assert correctness.get(q.question_id, False) is False

    def test_ratio_stratified_pool_insufficient(self) -> None:
        questions = _make_questions(10)
        correctness = {f"T-{i}": True for i in range(10)}
        with pytest.raises(ValueError, match="分层不足"):
            stratified_sample(
                questions=questions,
                correctness=correctness,
                size=10,
                correct_ratio=0.5,
                task_types=None,
                seed=42,
                min_per_class=None,
            )

    def test_task_types_filter(self) -> None:
        q_a = _make_questions(10, task_type="TypeA")
        q_b = _make_questions(10, task_type="TypeB")
        result = stratified_sample(
            questions=q_a + q_b,
            correctness={},
            size=5,
            correct_ratio=None,
            task_types=["TypeA"],
            seed=42,
            min_per_class=None,
        )
        assert all(q.task_type == "TypeA" for q in result)

    def test_unknown_correctness_treated_as_wrong(self) -> None:
        questions = _make_questions(20)
        correctness = {f"T-{i}": True for i in range(10)}
        result = stratified_sample(
            questions=questions,
            correctness=correctness,
            size=10,
            correct_ratio=0.5,
            task_types=None,
            seed=42,
            min_per_class=None,
        )
        n_correct = round(10 * 0.5)
        for q in result[:n_correct]:
            assert q.question_id in correctness

    def test_seed_reproducibility(self) -> None:
        questions = _make_questions(20)
        r1 = stratified_sample(
            questions=questions,
            correctness={},
            size=10,
            correct_ratio=None,
            task_types=None,
            seed=123,
            min_per_class=None,
        )
        r2 = stratified_sample(
            questions=questions,
            correctness={},
            size=10,
            correct_ratio=None,
            task_types=None,
            seed=123,
            min_per_class=None,
        )
        assert [q.question_id for q in r1] == [q.question_id for q in r2]

    def test_different_seeds_differ(self) -> None:
        questions = _make_questions(20)
        r1 = stratified_sample(
            questions=questions,
            correctness={},
            size=10,
            correct_ratio=None,
            task_types=None,
            seed=1,
            min_per_class=None,
        )
        r2 = stratified_sample(
            questions=questions,
            correctness={},
            size=10,
            correct_ratio=None,
            task_types=None,
            seed=2,
            min_per_class=None,
        )
        assert [q.question_id for q in r1] != [q.question_id for q in r2]

    def test_min_per_class_backfill(self) -> None:
        q_a = _make_questions(10, task_type="TypeA")
        q_b = _make_questions(10, task_type="TypeB")
        all_q = q_a + q_b
        correctness = {q.question_id: True for q in q_a[:5]}
        result = stratified_sample(
            questions=all_q,
            correctness=correctness,
            size=3,
            correct_ratio=None,
            task_types=None,
            seed=42,
            min_per_class=2,
        )
        type_counts: dict[str, int] = {}
        for q in result:
            type_counts[q.task_type] = type_counts.get(q.task_type, 0) + 1
        assert type_counts.get("TypeA", 0) >= 2
        assert type_counts.get("TypeB", 0) >= 2

    def test_min_per_class_partial_backfill(self) -> None:
        q_sparse = _make_questions(1, task_type="Sparse")
        q_main = _make_questions(10, task_type="Main")
        result = stratified_sample(
            questions=q_sparse + q_main,
            correctness={},
            size=5,
            correct_ratio=None,
            task_types=None,
            seed=42,
            min_per_class=3,
        )
        sparse_in_result = [q for q in result if q.task_type == "Sparse"]
        assert len(sparse_in_result) == 1

    def test_min_per_class_no_duplicates(self) -> None:
        q_a = _make_questions(5, task_type="TypeA")
        q_b = _make_questions(5, task_type="TypeB")
        result = stratified_sample(
            questions=q_a + q_b,
            correctness={},
            size=3,
            correct_ratio=None,
            task_types=None,
            seed=42,
            min_per_class=2,
        )
        ids = [q.question_id for q in result]
        assert len(ids) == len(set(ids))

    def test_backfill_enumerates_all_pool_types(self) -> None:
        q_main = _make_questions(10, task_type="Main")
        q_rare = _make_questions(3, task_type="Rare")
        result = stratified_sample(
            questions=q_main + q_rare,
            correctness={},
            size=2,
            correct_ratio=None,
            task_types=None,
            seed=0,
            min_per_class=1,
        )
        types_in_result = {q.task_type for q in result}
        assert "Rare" in types_in_result