Files
2026-07-14 05:55:08 -04:00

410 lines
14 KiB
Python

"""v2 生成器单元测试 — 验证 prompt 构建、响应解析、端到端生成。"""
from __future__ import annotations
import json
from typing import TYPE_CHECKING, Any
import pytest
from app.question_gen.families import RETRIEVAL_FAMILY, VISUAL_FAMILY
if TYPE_CHECKING:
from pathlib import Path
from app.question_gen.sampler_v2 import AnchorContext, MaterialContext
from app.tree.index import (
IndexMeta,
L1Card,
L1Node,
L2Card,
L2Node,
L3Card,
L3Node,
TreeIndex,
)
from core.types import LLMResponse
# ---------------------------------------------------------------------------
# Fixtures & Helpers
# ---------------------------------------------------------------------------
def _make_material(
*,
with_frames: bool = False,
with_cross_l2: bool = False,
) -> MaterialContext:
"""构造测试用 MaterialContext。"""
anchor = AnchorContext(node_id="L2_001", level=2, l2_id="L2_001")
frame_paths = ["/data/frames/f001.jpg", "/data/frames/f002.jpg"] if with_frames else []
cross_l2 = ["Person enters room and sits down."] if with_cross_l2 else []
return MaterialContext(
anchor=anchor,
source_nodes=("L2_001", "L3_001", "L3_002"),
subtitle_sentences=["He picks up the book.", "Then he starts reading."],
frame_paths=frame_paths,
cross_l2_texts=cross_l2,
)
def _make_vlm_response(content: str) -> LLMResponse:
"""构造 VLM 正常返回。"""
return LLMResponse(
content=content,
thinking="",
model="mock-vlm",
provider="mock",
prompt_tokens=50,
completion_tokens=30,
latency_ms=200,
ttft_ms=None,
max_inter_token_ms=None,
cache_hit=False,
call_id="mock-vlm-001",
)
class MockVLM:
"""可配置的 VLM mock — 记录调用并返回预设响应。"""
def __init__(self, response: LLMResponse) -> None:
self._response = response
self.calls: list[dict[str, Any]] = []
async def chat_with_images(
self,
messages: list[dict[str, Any]],
images: list[str | Path],
*,
session_id: str | None = None,
parent_call_id: str | None = None,
) -> LLMResponse:
"""记录调用并返回预设响应。"""
self.calls.append(
{
"messages": messages,
"images": images,
"session_id": session_id,
}
)
return self._response
def _make_tree() -> TreeIndex:
"""构造最小测试树。"""
l3_card = L3Card(
frame_summary="Person picks up a book from the shelf.",
visible_entities=["person", "book", "shelf"],
ongoing_actions=["picking up"],
visible_text=[],
spatial_layout="person in center",
visual_attributes={},
subtitle="He picks up the book.",
)
l3 = L3Node(id="L3_001", card=l3_card, frame_path="/data/frames/f001.jpg")
l3_card2 = L3Card(
frame_summary="Person starts reading the book.",
visible_entities=["person", "book"],
ongoing_actions=["reading"],
visible_text=[],
spatial_layout="person sitting",
visual_attributes={},
subtitle="Then he starts reading.",
)
l3_2 = L3Node(id="L3_002", card=l3_card2, frame_path="/data/frames/f002.jpg")
l2_card = L2Card(
event_description="A person picks up a book and starts reading.",
entities=["person", "book"],
actions=["pick up", "read"],
action_subjects=["person"],
visible_text=[],
spatial_relations="near shelf",
state_changes="book picked up",
subtitle="He picks up the book and reads.",
)
l2 = L2Node(id="L2_001", card=l2_card, children=[l3, l3_2])
l1_card = L1Card(
scene_summary="Library scene.",
main_setting="library",
key_entities=["person", "book"],
main_actions=["reading"],
topic_keywords=["library"],
visible_text=[],
temporal_flow="enter → read",
)
l1 = L1Node(id="L1_001", card=l1_card, children=[l2])
meta = IndexMeta(source_path="test.mp4", modality="video")
return TreeIndex(metadata=meta, roots=[l1])
_VALID_VLM_OUTPUT = json.dumps(
{
"question": "What does the person do after picking up the book?",
"options": ["A. Reads it", "B. Puts it back", "C. Throws it", "D. Burns it"],
"answer": "A",
"difficulty": "medium",
}
)
# ---------------------------------------------------------------------------
# TestBuildV2Prompt
# ---------------------------------------------------------------------------
class TestBuildV2Prompt:
"""验证 _build_v2_prompt 的 prompt 构造逻辑。"""
def test_includes_family_template(self) -> None:
"""prompt 中应包含家族模板的核心指令。"""
from app.question_gen.generator_v2 import _build_v2_prompt
material = _make_material()
messages, frame_paths = _build_v2_prompt(
prompt_template=RETRIEVAL_FAMILY.prompt_template,
strategy_name=RETRIEVAL_FAMILY.name,
material=material,
task_type="Action Reasoning",
seq=1,
)
# messages 至少有 system + user
assert len(messages) >= 2
# system message 中应包含 retrieval 家族相关指令
system_content = messages[0]["content"]
assert "retrieval" in system_content.lower() or "factual" in system_content.lower()
def test_reject_reason_injected(self) -> None:
"""当 reject_reason 不为 None 时,应注入到 prompt 中。"""
from app.question_gen.generator_v2 import _build_v2_prompt
material = _make_material()
reject_msg = "The question leaked information from the answer options."
messages, _ = _build_v2_prompt(
prompt_template=RETRIEVAL_FAMILY.prompt_template,
strategy_name=RETRIEVAL_FAMILY.name,
material=material,
task_type="Action Reasoning",
seq=2,
reject_reason=reject_msg,
)
# reject_reason 应出现在某个 message 内容中
all_content = " ".join(m["content"] for m in messages)
assert reject_msg in all_content
def test_frame_paths_from_material(self) -> None:
"""返回的 frame_paths 应来自 material.frame_paths。"""
from app.question_gen.generator_v2 import _build_v2_prompt
material = _make_material(with_frames=True)
_, frame_paths = _build_v2_prompt(
prompt_template=VISUAL_FAMILY.prompt_template,
strategy_name=VISUAL_FAMILY.name,
material=material,
task_type="Object Recognition",
seq=1,
)
assert frame_paths == ["/data/frames/f001.jpg", "/data/frames/f002.jpg"]
def test_no_frames_returns_empty(self) -> None:
"""无帧素材时 frame_paths 应为空列表。"""
from app.question_gen.generator_v2 import _build_v2_prompt
material = _make_material(with_frames=False)
_, frame_paths = _build_v2_prompt(
prompt_template=RETRIEVAL_FAMILY.prompt_template,
strategy_name=RETRIEVAL_FAMILY.name,
material=material,
task_type="Action Reasoning",
seq=1,
)
assert frame_paths == []
# ---------------------------------------------------------------------------
# TestParseV2Response
# ---------------------------------------------------------------------------
class TestParseV2Response:
"""验证 _parse_v2_response 的解析与校验逻辑。"""
def test_valid_json(self) -> None:
"""正确 JSON → 生成 CandidateQuestion。"""
from app.question_gen.generator_v2 import CandidateQuestion, _parse_v2_response
result = _parse_v2_response(
raw=_VALID_VLM_OUTPUT,
video_id="v-001",
task_type="Action Reasoning",
skill_target="M1",
seq=3,
source_nodes=("L2_001", "L3_001"),
)
assert isinstance(result, CandidateQuestion)
assert result.question_id == "v-001_Action Reasoning_0003"
assert result.question == "What does the person do after picking up the book?"
assert result.options == ("A. Reads it", "B. Puts it back", "C. Throws it", "D. Burns it")
assert result.answer == "A"
assert result.difficulty == "medium"
assert result.video_id == "v-001"
assert result.task_type == "Action Reasoning"
assert result.skill_target == "M1"
assert result.source_nodes == ("L2_001", "L3_001")
def test_missing_field_raises(self) -> None:
"""缺少必填字段 → 抛出 ValueError。"""
from app.question_gen.generator_v2 import _parse_v2_response
incomplete = json.dumps(
{
"question": "What happens?",
"options": ["A. X", "B. Y"],
# 缺少 answer 和 difficulty
}
)
with pytest.raises(ValueError, match="answer"):
_parse_v2_response(
raw=incomplete,
video_id="v-001",
task_type="Action Reasoning",
skill_target="M1",
seq=1,
source_nodes=("L2_001",),
)
def test_invalid_answer_raises(self) -> None:
"""answer 不在 A-D 范围内 → 抛出 ValueError。"""
from app.question_gen.generator_v2 import _parse_v2_response
bad_answer = json.dumps(
{
"question": "What happens?",
"options": ["A. X", "B. Y", "C. Z", "D. W"],
"answer": "E",
"difficulty": "easy",
}
)
with pytest.raises(ValueError, match="answer"):
_parse_v2_response(
raw=bad_answer,
video_id="v-001",
task_type="Action Reasoning",
skill_target="M1",
seq=1,
source_nodes=("L2_001",),
)
def test_json_repair_handles_trailing_comma(self) -> None:
"""json_repair 应能修复常见 JSON 错误。"""
from app.question_gen.generator_v2 import _parse_v2_response
malformed = '{"question": "What?", "options": ["A. X", "B. Y", "C. Z", "D. W",], "answer": "B", "difficulty": "easy",}'
result = _parse_v2_response(
raw=malformed,
video_id="v-001",
task_type="Action Reasoning",
skill_target="M1",
seq=1,
source_nodes=("L2_001",),
)
assert result.answer == "B"
def test_markdown_wrapped_json(self) -> None:
"""被 markdown 代码块包裹的 JSON 应正常解析。"""
from app.question_gen.generator_v2 import _parse_v2_response
wrapped = f"```json\n{_VALID_VLM_OUTPUT}\n```"
result = _parse_v2_response(
raw=wrapped,
video_id="v-001",
task_type="Action Reasoning",
skill_target="M1",
seq=5,
source_nodes=("L2_001",),
)
assert result.question_id == "v-001_Action Reasoning_0005"
# ---------------------------------------------------------------------------
# TestGenerateOneV2
# ---------------------------------------------------------------------------
class TestGenerateOneV2:
"""验证 generate_one_v2 端到端流程。"""
@pytest.mark.asyncio()
async def test_happy_path(self) -> None:
"""正常流程:VLM 返回有效 JSON → 生成 CandidateQuestion。"""
from app.question_gen.generator_v2 import CandidateQuestion, generate_one_v2
mock_vlm = MockVLM(_make_vlm_response(_VALID_VLM_OUTPUT))
tree = _make_tree()
material = _make_material(with_frames=True)
result = await generate_one_v2(
vlm=mock_vlm,
tree=tree,
material=material,
task_type="Action Reasoning",
seq=1,
video_id="v-001",
prompt_template=RETRIEVAL_FAMILY.prompt_template,
strategy_name=RETRIEVAL_FAMILY.name,
skill_target=RETRIEVAL_FAMILY.skill_target,
session_id="session-test-001",
)
assert isinstance(result, CandidateQuestion)
assert result.question_id == "v-001_Action Reasoning_0001"
assert result.video_id == "v-001"
assert result.skill_target == "M1"
assert result.source_nodes == ("L2_001", "L3_001", "L3_002")
assert result.subtitle_sentences == ("He picks up the book.", "Then he starts reading.")
assert result.frame_paths == ("/data/frames/f001.jpg", "/data/frames/f002.jpg")
# VLM 应被调用一次
assert len(mock_vlm.calls) == 1
assert mock_vlm.calls[0]["session_id"] == "session-test-001"
@pytest.mark.asyncio()
async def test_reject_reason_forwarded(self) -> None:
"""reject_reason 应被传入 prompt 构建。"""
from app.question_gen.generator_v2 import generate_one_v2
mock_vlm = MockVLM(_make_vlm_response(_VALID_VLM_OUTPUT))
tree = _make_tree()
material = _make_material()
await generate_one_v2(
vlm=mock_vlm,
tree=tree,
material=material,
task_type="Action Reasoning",
seq=2,
video_id="v-001",
prompt_template=RETRIEVAL_FAMILY.prompt_template,
strategy_name=RETRIEVAL_FAMILY.name,
skill_target=RETRIEVAL_FAMILY.skill_target,
reject_reason="Answer was too obvious",
session_id="session-test-002",
)
# 验证 VLM 调用的 messages 中包含 reject_reason
call = mock_vlm.calls[0]
all_content = " ".join(m["content"] for m in call["messages"])
assert "Answer was too obvious" in all_content