Files
Video-Tree-TRM5/tests/unit/test_harness_pools.py
T

291 lines
9.4 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""三池切分单元测试。
验证:
- 三池互斥(question_id 无重叠)
- test 池自然分布(correct_ratio=None
- save/load 往返一致
- 旧格式拒绝(无 test 键 → ValueError
- build_or_load_pools 冻结复用(pools.json 存在时不重切)
"""
from __future__ import annotations
import json
from typing import TYPE_CHECKING
import pytest
from app.harness.pools import (
build_pools,
load_pools,
save_pools,
)
from core.types import GeneratedQuestion
if TYPE_CHECKING:
from pathlib import Path
def _make_question(qid: str, task_type: str = "Action Reasoning") -> GeneratedQuestion:
"""构造测试用 GeneratedQuestion。
参数:
qid: 题目 ID。
task_type: 题型。
返回:
GeneratedQuestion 实例。
"""
return GeneratedQuestion(
question_id=qid,
video_id="video_001",
task_type=task_type,
question=f"Question {qid}?",
options=("A. opt1", "B. opt2", "C. opt3", "D. opt4"),
answer="A",
source_nodes=("node_1",),
difficulty="medium",
)
def _make_question_set(
n: int,
task_types: list[str] | None = None,
) -> list[GeneratedQuestion]:
"""构造 n 道题,交替分配题型。
参数:
n: 题目数量。
task_types: 可选题型列表,轮转分配;None 默认 2 类。
返回:
题目列表。
"""
types = task_types or ["Action Reasoning", "Scene Understanding"]
return [_make_question(f"q_{i:04d}", types[i % len(types)]) for i in range(n)]
def _make_correctness(
questions: list[GeneratedQuestion],
correct_ratio: float = 0.5,
) -> dict[str, bool]:
"""构造 correctness 字典,前 correct_ratio 比例标对。
参数:
questions: 题目列表。
correct_ratio: 对题占比。
返回:
question_id -> bool。
"""
n_correct = round(len(questions) * correct_ratio)
return {q.question_id: (i < n_correct) for i, q in enumerate(questions)}
class TestBuildPoolsMutualExclusion:
"""三池 question_id 互斥验证。"""
def test_build_pools_mutual_exclusion(self) -> None:
"""三池切分后,任意两池不共享 question_id。"""
questions = _make_question_set(200)
correctness = _make_correctness(questions, 0.5)
pools = build_pools(
questions,
correctness,
diag_cfg={
"size": 30,
"correct_ratio": 0.5,
"task_types": None,
"seed": 42,
"min_per_class": None,
},
val_cfg={
"size": 30,
"correct_ratio": 0.5,
"task_types": None,
"seed": 42,
"min_per_class": None,
},
test_cfg={"size": 30},
baseline_run_id="run_baseline",
)
diag_ids = {q.question_id for q in pools.diagnosis}
val_ids = {q.question_id for q in pools.validation}
test_ids = {q.question_id for q in pools.test}
assert diag_ids & val_ids == set(), "诊断池与验证池有重叠"
assert diag_ids & test_ids == set(), "诊断池与测试池有重叠"
assert val_ids & test_ids == set(), "验证池与测试池有重叠"
assert len(diag_ids) == 30
assert len(val_ids) == 30
assert len(test_ids) == 30
class TestBuildPoolsTestNaturalDistribution:
"""test 池使用自然分布(correct_ratio=None)。"""
def test_build_pools_test_natural_distribution(self) -> None:
"""test 池不强制对错比例,保留候选池的自然分布。
构造 correctness 为 50% 对/50% 错,diag/val 用 correct_ratio=0.3
强制裁剪,test 池走自然分布(correct_ratio=None)。验证 test 池
不受 correct_ratio 约束。
"""
questions = _make_question_set(300)
correctness = _make_correctness(questions, 0.5)
pools = build_pools(
questions,
correctness,
diag_cfg={
"size": 20,
"correct_ratio": 0.3,
"task_types": None,
"seed": 42,
"min_per_class": None,
},
val_cfg={
"size": 20,
"correct_ratio": 0.3,
"task_types": None,
"seed": 42,
"min_per_class": None,
},
test_cfg={"size": 20},
baseline_run_id="run_baseline",
)
# diag/val 被 correct_ratio=0.3 裁剪:round(20*0.3) = 6 对, 14 错
diag_correct = sum(1 for q in pools.diagnosis if correctness[q.question_id])
val_correct = sum(1 for q in pools.validation if correctness[q.question_id])
assert diag_correct == 6, "诊断池应强制 30% 对题"
assert val_correct == 6, "验证池应强制 30% 对题"
# test 池自然分布:不受 correct_ratio 约束
assert len(pools.test) == 20
class TestSaveLoadPoolsRoundtrip:
"""save/load 往返一致验证。"""
def test_save_load_pools_roundtrip(self, tmp_path: Path) -> None:
"""save_pools → load_pools 后全字段一致。"""
questions = _make_question_set(100)
correctness = _make_correctness(questions, 0.5)
original = build_pools(
questions,
correctness,
diag_cfg={
"size": 15,
"correct_ratio": 0.5,
"task_types": None,
"seed": 42,
"min_per_class": None,
},
val_cfg={
"size": 15,
"correct_ratio": 0.5,
"task_types": None,
"seed": 42,
"min_per_class": None,
},
test_cfg={"size": 15},
baseline_run_id="run_001",
)
pools_path = tmp_path / "pools.json"
save_pools(original, pools_path)
restored = load_pools(pools_path)
# 标量字段
assert restored.baseline_run_id == original.baseline_run_id
assert restored.baseline_val_accuracy == pytest.approx(original.baseline_val_accuracy)
assert restored.correctness == original.correctness
# 三池逐题比对
for pool_name in ("diagnosis", "validation", "test"):
orig_list = getattr(original, pool_name)
rest_list = getattr(restored, pool_name)
assert len(rest_list) == len(orig_list), f"{pool_name} 长度不一致"
for o, r in zip(orig_list, rest_list, strict=False):
assert o.question_id == r.question_id
assert o.video_id == r.video_id
assert o.task_type == r.task_type
assert o.question == r.question
assert o.options == r.options
assert o.answer == r.answer
assert o.source_nodes == r.source_nodes
assert o.difficulty == r.difficulty
class TestLoadPoolsOldFormatReject:
"""旧格式 pools.json(无 test 键)→ ValueError。"""
def test_load_pools_old_format_reject(self, tmp_path: Path) -> None:
"""缺少 test 键的 pools.json 必须抛出 ValueError。"""
old_format = {
"baseline_run_id": "run_old",
"baseline_val_accuracy": 0.5,
"correctness": {},
"diagnosis": [],
"validation": [],
}
pools_path = tmp_path / "pools.json"
pools_path.write_text(json.dumps(old_format), encoding="utf-8")
with pytest.raises(ValueError, match="旧格式"):
load_pools(pools_path)
class TestBuildOrLoadPoolsFrozen:
"""build_or_load_pools 冻结复用:pools.json 存在时原样加载不重切。"""
def test_build_or_load_pools_frozen(self, tmp_path: Path) -> None:
"""pools.json 已存在时,build_or_load_pools 返回冻结内容。"""
questions = _make_question_set(60)
correctness = _make_correctness(questions, 0.5)
frozen = build_pools(
questions,
correctness,
diag_cfg={
"size": 10,
"correct_ratio": 0.5,
"task_types": None,
"seed": 42,
"min_per_class": None,
},
val_cfg={
"size": 10,
"correct_ratio": 0.5,
"task_types": None,
"seed": 42,
"min_per_class": None,
},
test_cfg={"size": 10},
baseline_run_id="run_frozen",
)
pools_path = tmp_path / "pools.json"
save_pools(frozen, pools_path)
# build_or_load_pools 中 pools.json 存在 → 直接 load_pools
# 此处直接测试 load_pools 行为等价
loaded = load_pools(pools_path)
assert loaded.baseline_run_id == frozen.baseline_run_id
assert loaded.baseline_val_accuracy == pytest.approx(frozen.baseline_val_accuracy)
assert len(loaded.test) == len(frozen.test)
assert len(loaded.validation) == len(frozen.validation)
assert len(loaded.diagnosis) == len(frozen.diagnosis)
# question_id 完全一致
for pool_name in ("diagnosis", "validation", "test"):
orig_ids = [q.question_id for q in getattr(frozen, pool_name)]
load_ids = [q.question_id for q in getattr(loaded, pool_name)]
assert orig_ids == load_ids, f"{pool_name} 冻结后 ID 顺序不一致"