Files
Video-Tree-TRM5/tests/unit/test_strategy.py
T
2026-07-14 06:51:50 -04:00

143 lines
5.1 KiB
Python

"""TaskTypeStrategy Protocol 与 BaseTaskTypeStrategy 单元测试。"""
from __future__ import annotations
import random
import pytest
from app.question_gen.families import (
ENUMERATION_FAMILY,
REASONING_FAMILY,
RETRIEVAL_FAMILY,
SPATIAL_FAMILY,
VISUAL_FAMILY,
)
from app.question_gen.strategy import (
BaseTaskTypeStrategy,
SubPattern,
get_strategy,
register_strategy,
)
class TestBaseTaskTypeStrategy:
"""BaseTaskTypeStrategy 封装 family 行为。"""
def test_task_type_and_strategy_name(self):
"""task_type 和 strategy_name 正确返回。"""
s = BaseTaskTypeStrategy(task_type="Object Recognition", family=RETRIEVAL_FAMILY)
assert s.task_type == "Object Recognition"
assert s.strategy_name == "RETRIEVAL"
def test_sampling_from_family(self):
"""采样参数从绑定的 family 读取。"""
s = BaseTaskTypeStrategy(task_type="Temporal Reasoning", family=ENUMERATION_FAMILY)
assert s.sampling_level == 1
assert s.sampling_constraint.min_l3_nodes == 5
def test_skill_target_from_family(self):
"""skill_target 从 family 读取。"""
s = BaseTaskTypeStrategy(task_type="Action Reasoning", family=REASONING_FAMILY)
assert s.skill_target == "M2"
def test_leak_probe_template(self):
"""leak_probe_template 从 family.leak_profile 读取。"""
s = BaseTaskTypeStrategy(task_type="Spatial Reasoning", family=SPATIAL_FAMILY)
assert s.leak_probe_template == "gate_leak_spatial.md"
def test_prompt_template_from_family(self):
"""prompt_template 从 family 读取。"""
s = BaseTaskTypeStrategy(task_type="OCR Problems", family=VISUAL_FAMILY)
assert s.prompt_template == "visual.md"
def test_select_sub_pattern_returns_none(self):
"""BaseTaskTypeStrategy 无子模式。"""
s = BaseTaskTypeStrategy(task_type="Object Recognition", family=RETRIEVAL_FAMILY)
rng = random.Random(42)
assert s.select_sub_pattern(rng) is None
def test_extra_gates_returns_empty(self):
"""BaseTaskTypeStrategy 无额外 gate。"""
s = BaseTaskTypeStrategy(task_type="Object Recognition", family=RETRIEVAL_FAMILY)
assert s.extra_gates(None) == []
class TestSamplingLevelMapping:
"""BaseTaskTypeStrategy 的 sampling_level 从 _TASK_TYPE_TO_LEVEL 读取。"""
def test_l3_types(self):
"""L3 题型。"""
for tt in ("Object Recognition",):
s = BaseTaskTypeStrategy(task_type=tt, family=RETRIEVAL_FAMILY)
assert s.sampling_level == 3, f"{tt} should be L3"
def test_l2_types(self):
"""L2 题型。"""
s = BaseTaskTypeStrategy(task_type="Action Reasoning", family=REASONING_FAMILY)
assert s.sampling_level == 2
def test_l1_types(self):
"""L1 题型。"""
s = BaseTaskTypeStrategy(task_type="Temporal Reasoning", family=ENUMERATION_FAMILY)
assert s.sampling_level == 1
class TestStrategyRegistry:
"""注册表查找。"""
def test_get_unregistered_returns_base(self):
"""未注册题型返回 BaseTaskTypeStrategy。"""
s = get_strategy("Spatial Reasoning")
assert isinstance(s, BaseTaskTypeStrategy)
assert s.task_type == "Spatial Reasoning"
def test_register_and_get(self):
"""注册后 get 返回注册的策略。"""
custom = BaseTaskTypeStrategy(task_type="Object Recognition", family=RETRIEVAL_FAMILY)
register_strategy(custom)
assert get_strategy("Object Recognition") is custom
class TestSubPattern:
"""SubPattern 数据类。"""
def test_frozen(self):
"""SubPattern 不可变。"""
sp = SubPattern(
name="test",
weight=0.5,
sampling_level_override=None,
constraint_override=None,
instruction="test instruction",
positive_examples=[],
negative_examples=[],
distractor_rules="",
)
with pytest.raises(AttributeError):
sp.name = "changed"
class TestActionRecognitionRegistration:
"""AR 策略注册后 get_strategy 返回特化实例。"""
def test_get_strategy_returns_ar_strategy(self):
"""get_strategy('Action Recognition') 返回 ActionRecognitionStrategy。"""
from app.question_gen.strategy_action_recognition import ActionRecognitionStrategy
s = get_strategy("Action Recognition")
assert isinstance(s, ActionRecognitionStrategy)
assert s.task_type == "Action Recognition"
assert s.strategy_name == "ACTION_RECOGNITION"
def test_ar_not_base_strategy(self):
"""get_strategy('Action Recognition') 不再返回 BaseTaskTypeStrategy。"""
s = get_strategy("Action Recognition")
assert not isinstance(s, BaseTaskTypeStrategy)
def test_other_types_still_base(self):
"""其他题型仍返回 BaseTaskTypeStrategy。"""
for tt in ("Object Recognition", "Temporal Reasoning", "Spatial Reasoning"):
s = get_strategy(tt)
assert isinstance(s, BaseTaskTypeStrategy), f"{tt} 应该是 BaseTaskTypeStrategy"