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

660 lines
21 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.
"""app/harness/inference.py 单元测试。
测试覆盖:
- run_inference 基本流程(mock LLM + tool_dispatch
- 异常时 prediction 仍落库(stop_reason=error
- _to_text_field 归一化
- run_id 空串 → ValueError
- _aggregate_results 内存聚合
- 空 questions 列表零值返回
- 并发控制 Semaphore
- plugins_factory 调用
"""
from __future__ import annotations
import json
from typing import Any
from unittest.mock import AsyncMock
import pytest
from app.harness.inference import (
InferenceResult,
_aggregate_results,
_to_text_field,
run_inference,
)
from app.harness.log import HarnessLog
from core.types import GeneratedQuestion, LLMResponse
# ── 测试基础设施 ──────────────────────────────────────────────────
def _make_question(
question_id: str = "q1",
video_id: str = "v1",
task_type: str = "Action Reasoning",
answer: str = "B",
) -> GeneratedQuestion:
"""构造测试用题目。"""
return GeneratedQuestion(
question_id=question_id,
video_id=video_id,
task_type=task_type,
question="测试问题",
options=("A. 选项A", "B. 选项B", "C. 选项C", "D. 选项D"),
answer=answer,
source_nodes=("L1_001",),
difficulty="medium",
)
def _make_llm_response(answer: str = "B") -> LLMResponse:
"""构造测试用 LLMResponsesubmit_answer 场景)。"""
content = json.dumps(
{
"reflect": {"observation": "找到答案"},
"plan": {"next_step": "提交"},
"action": {
"tool": "submit_answer",
"args": {
"answer": answer,
"evidence": "证据文本",
"reasoning": "推理过程",
},
},
}
)
return LLMResponse(
content=content,
thinking="思考过程",
model="test-model",
provider="test",
prompt_tokens=100,
completion_tokens=50,
latency_ms=200,
ttft_ms=30.0,
max_inter_token_ms=5.0,
cache_hit=False,
call_id="test-call-001",
)
def _make_error_llm_response() -> LLMResponse:
"""构造触发解析失败的 LLMResponse。"""
return LLMResponse(
content="这不是JSON",
thinking="",
model="test-model",
provider="test",
prompt_tokens=10,
completion_tokens=5,
latency_ms=50,
ttft_ms=10.0,
max_inter_token_ms=2.0,
cache_hit=False,
call_id="test-call-err",
)
async def _stub_tool_dispatch(
tool_name: str, args: dict[str, Any], *, context: dict[str, Any]
) -> str:
"""测试用工具调度函数。"""
if tool_name == "submit_answer":
return "答案已提交"
raise ValueError(f"未知工具: {tool_name}")
def _stub_prompt_builder(qa: GeneratedQuestion) -> tuple[str, str]:
"""测试用 prompt 构建函数。"""
return "系统提示词", f"用户问题: {qa.question}"
@pytest.fixture
def harness_log(tmp_path: Any, request: Any) -> HarnessLog:
"""创建临时 HarnessLog 实例。
使用 test 节点名称的 hash 作为 db 文件名,避免冲突。
run_id 固定为 "test-run",实际 run_inference 中传入的 run_id
由 HarnessLog.insert 自动覆盖为 HarnessLog 构造时的值。
"""
db_name = f"harness_{id(request)}.db"
db_path = str(tmp_path / db_name)
log = HarnessLog(db_path, "test-run")
yield log
log.close()
# ── 测试用例 ──────────────────────────────────────────────────
class TestToTextField:
"""_to_text_field 归一化测试。"""
@pytest.mark.asyncio
async def test_string_passthrough(self) -> None:
"""字符串原样返回。"""
assert _to_text_field("hello") == "hello"
@pytest.mark.asyncio
async def test_empty_string(self) -> None:
"""空字符串原样返回。"""
assert _to_text_field("") == ""
@pytest.mark.asyncio
async def test_list_serialized(self) -> None:
"""list 被 JSON 序列化。"""
result = _to_text_field(["a", "b"])
assert result == '["a", "b"]'
@pytest.mark.asyncio
async def test_dict_serialized(self) -> None:
"""dict 被 JSON 序列化。"""
result = _to_text_field({"key": "值"})
assert '"key"' in result
assert '"值"' in result
@pytest.mark.asyncio
async def test_int_serialized(self) -> None:
"""int 被 JSON 序列化。"""
assert _to_text_field(42) == "42"
@pytest.mark.asyncio
async def test_none_serialized(self) -> None:
"""None 被 JSON 序列化。"""
assert _to_text_field(None) == "null"
@pytest.mark.asyncio
async def test_unicode_preserved(self) -> None:
"""ensure_ascii=False 保留中文。"""
result = _to_text_field(["中文"])
assert "中文" in result
assert "\\u" not in result
class TestAggregateResults:
"""_aggregate_results 内存聚合测试。"""
@pytest.mark.asyncio
async def test_empty_records(self) -> None:
"""空列表返回零值 InferenceResult。"""
result = _aggregate_results([], "run-empty")
assert result.run_id == "run-empty"
assert result.accuracy == 0.0
assert result.total == 0
assert result.correct == 0
assert result.per_task_type == {}
assert result.steps_mean == 0.0
assert result.token_usage == {"prompt_tokens": 0, "completion_tokens": 0}
assert result.stop_reason_counts == {}
@pytest.mark.asyncio
async def test_single_correct(self) -> None:
"""单条正确记录 → accuracy=1.0。"""
records = [
{
"prediction": "B",
"answer": "B",
"task_type": "AR",
"steps_used": 3,
"prompt_tokens": 100,
"completion_tokens": 50,
"stop_reason": "finished",
}
]
result = _aggregate_results(records, "run-1")
assert result.accuracy == 1.0
assert result.total == 1
assert result.correct == 1
assert result.steps_mean == 3.0
@pytest.mark.asyncio
async def test_mixed_correct_wrong(self) -> None:
"""混合正确/错误 → 准确率与步数均正确聚合。"""
records = [
{
"prediction": "B",
"answer": "B",
"task_type": "AR",
"steps_used": 2,
"prompt_tokens": 100,
"completion_tokens": 50,
"stop_reason": "finished",
},
{
"prediction": "C",
"answer": "A",
"task_type": "AR",
"steps_used": 4,
"prompt_tokens": 200,
"completion_tokens": 100,
"stop_reason": "budget_exceeded",
},
{
"prediction": "D",
"answer": "D",
"task_type": "SP",
"steps_used": 1,
"prompt_tokens": 50,
"completion_tokens": 25,
"stop_reason": "finished",
},
]
result = _aggregate_results(records, "run-mix")
assert result.total == 3
assert result.correct == 2
assert abs(result.accuracy - 2 / 3) < 1e-9
assert abs(result.steps_mean - 7 / 3) < 1e-9
assert result.token_usage == {"prompt_tokens": 350, "completion_tokens": 175}
assert result.stop_reason_counts == {"finished": 2, "budget_exceeded": 1}
@pytest.mark.asyncio
async def test_per_task_type_grouping(self) -> None:
"""按 task_type 分组聚合。"""
records = [
{
"prediction": "B",
"answer": "B",
"task_type": "AR",
"steps_used": 1,
"prompt_tokens": 10,
"completion_tokens": 5,
"stop_reason": "finished",
},
{
"prediction": "A",
"answer": "C",
"task_type": "AR",
"steps_used": 2,
"prompt_tokens": 20,
"completion_tokens": 10,
"stop_reason": "finished",
},
{
"prediction": "D",
"answer": "D",
"task_type": "SP",
"steps_used": 3,
"prompt_tokens": 30,
"completion_tokens": 15,
"stop_reason": "finished",
},
]
result = _aggregate_results(records, "run-task")
assert "AR" in result.per_task_type
assert "SP" in result.per_task_type
assert result.per_task_type["AR"]["total"] == 2
assert result.per_task_type["AR"]["correct"] == 1
assert result.per_task_type["AR"]["accuracy"] == 0.5
assert result.per_task_type["SP"]["total"] == 1
assert result.per_task_type["SP"]["correct"] == 1
assert result.per_task_type["SP"]["accuracy"] == 1.0
class TestRunIdValidation:
"""run_id 校验测试。"""
@pytest.mark.asyncio
async def test_empty_string_raises(self, harness_log: HarnessLog) -> None:
"""空串 run_id → ValueError。"""
llm = AsyncMock()
with pytest.raises(ValueError, match="run_id 不得为空"):
await run_inference(
[],
llm=llm,
tool_dispatch_fn=_stub_tool_dispatch,
prompt_builder=_stub_prompt_builder,
log=harness_log,
run_id="",
concurrency=1,
max_steps=10,
skill_mode="auto",
)
@pytest.mark.asyncio
async def test_whitespace_only_raises(self, harness_log: HarnessLog) -> None:
"""纯空白 run_id → ValueError。"""
llm = AsyncMock()
with pytest.raises(ValueError, match="run_id 不得为空"):
await run_inference(
[],
llm=llm,
tool_dispatch_fn=_stub_tool_dispatch,
prompt_builder=_stub_prompt_builder,
log=harness_log,
run_id=" ",
concurrency=1,
max_steps=10,
skill_mode="auto",
)
class TestEmptyQuestions:
"""空题目列表测试。"""
@pytest.mark.asyncio
async def test_empty_questions_returns_zero(self, harness_log: HarnessLog) -> None:
"""空 questions 列表直接返回零值 InferenceResult。"""
llm = AsyncMock()
result = await run_inference(
[],
llm=llm,
tool_dispatch_fn=_stub_tool_dispatch,
prompt_builder=_stub_prompt_builder,
log=harness_log,
run_id="run-empty",
concurrency=1,
max_steps=10,
skill_mode="auto",
)
assert isinstance(result, InferenceResult)
assert result.run_id == "run-empty"
assert result.accuracy == 0.0
assert result.total == 0
assert result.correct == 0
# LLM 未被调用
llm.chat.assert_not_called()
class TestRunInferenceBasic:
"""run_inference 基本流程测试。"""
@pytest.mark.asyncio
async def test_single_question_correct(self, harness_log: HarnessLog) -> None:
"""单题正确推理 → accuracy=1.0, stop_reason=finished。"""
llm = AsyncMock()
llm.chat.return_value = _make_llm_response(answer="B")
result = await run_inference(
[_make_question(answer="B")],
llm=llm,
tool_dispatch_fn=_stub_tool_dispatch,
prompt_builder=_stub_prompt_builder,
log=harness_log,
run_id="run-basic",
concurrency=1,
max_steps=10,
skill_mode="auto",
)
assert result.accuracy == 1.0
assert result.total == 1
assert result.correct == 1
assert result.stop_reason_counts.get("finished") == 1
@pytest.mark.asyncio
async def test_single_question_wrong(self, harness_log: HarnessLog) -> None:
"""单题错误推理 → accuracy=0.0。"""
llm = AsyncMock()
llm.chat.return_value = _make_llm_response(answer="C")
result = await run_inference(
[_make_question(answer="B")],
llm=llm,
tool_dispatch_fn=_stub_tool_dispatch,
prompt_builder=_stub_prompt_builder,
log=harness_log,
run_id="run-wrong",
concurrency=1,
max_steps=10,
skill_mode="auto",
)
assert result.accuracy == 0.0
assert result.total == 1
assert result.correct == 0
@pytest.mark.asyncio
async def test_multiple_questions_concurrent(self, harness_log: HarnessLog) -> None:
"""3 题并发推理 → 结果正确聚合。"""
llm = AsyncMock()
llm.chat.return_value = _make_llm_response(answer="B")
questions = [
_make_question(question_id="q1", answer="B"),
_make_question(question_id="q2", answer="B"),
_make_question(question_id="q3", answer="A"),
]
result = await run_inference(
questions,
llm=llm,
tool_dispatch_fn=_stub_tool_dispatch,
prompt_builder=_stub_prompt_builder,
log=harness_log,
run_id="run-multi",
concurrency=3,
max_steps=10,
skill_mode="auto",
)
assert result.total == 3
assert result.correct == 2
assert abs(result.accuracy - 2 / 3) < 1e-9
@pytest.mark.asyncio
async def test_token_usage_accumulated(self, harness_log: HarnessLog) -> None:
"""多题 token 累加验证。"""
llm = AsyncMock()
llm.chat.return_value = _make_llm_response(answer="B")
questions = [
_make_question(question_id="q1"),
_make_question(question_id="q2"),
]
result = await run_inference(
questions,
llm=llm,
tool_dispatch_fn=_stub_tool_dispatch,
prompt_builder=_stub_prompt_builder,
log=harness_log,
run_id="run-token",
concurrency=2,
max_steps=10,
skill_mode="auto",
)
assert result.token_usage["prompt_tokens"] == 200
assert result.token_usage["completion_tokens"] == 100
class TestPredictionAlwaysWritten:
"""异常时 prediction 仍落库测试。"""
@pytest.mark.asyncio
async def test_error_still_persisted(self, harness_log: HarnessLog) -> None:
"""LLM 调用异常时,prediction 仍以 stop_reason=error 落库。"""
llm = AsyncMock()
llm.chat.side_effect = RuntimeError("LLM API 不可用")
result = await run_inference(
[_make_question()],
llm=llm,
tool_dispatch_fn=_stub_tool_dispatch,
prompt_builder=_stub_prompt_builder,
log=harness_log,
run_id="run-error",
concurrency=1,
max_steps=10,
skill_mode="auto",
)
assert result.total == 1
assert result.correct == 0
assert result.stop_reason_counts.get("error") == 1
# 验证 DB 中的记录(HarnessLog.insert 使用构造时的 run_id
rows = harness_log.query("SELECT * FROM predictions WHERE run_id = ?", ("test-run",))
assert len(rows) == 1
assert rows[0]["stop_reason"] == "error"
assert rows[0]["prediction"] is None
@pytest.mark.asyncio
async def test_parse_error_still_persisted(self, harness_log: HarnessLog) -> None:
"""LLM 返回非 JSON 内容,parse_error 后 prediction 仍落库。"""
llm = AsyncMock()
llm.chat.return_value = _make_error_llm_response()
result = await run_inference(
[_make_question()],
llm=llm,
tool_dispatch_fn=_stub_tool_dispatch,
prompt_builder=_stub_prompt_builder,
log=harness_log,
run_id="run-parse-err",
concurrency=1,
max_steps=10,
skill_mode="auto",
)
assert result.total == 1
# HarnessLog.insert 使用构造时的 run_id
rows = harness_log.query("SELECT * FROM predictions WHERE run_id = ?", ("test-run",))
assert len(rows) == 1
assert rows[0]["prediction"] is None
class TestPluginsFactory:
"""plugins_factory 调用测试。"""
@pytest.mark.asyncio
async def test_factory_called_per_question(self, harness_log: HarnessLog) -> None:
"""每题调用 plugins_factory,传入 (video_id, question_id)。"""
llm = AsyncMock()
llm.chat.return_value = _make_llm_response(answer="B")
factory_calls: list[tuple[str, str]] = []
def _factory(video_id: str, question_id: str) -> list[object]:
factory_calls.append((video_id, question_id))
return []
questions = [
_make_question(question_id="q1", video_id="v1"),
_make_question(question_id="q2", video_id="v2"),
]
await run_inference(
questions,
llm=llm,
tool_dispatch_fn=_stub_tool_dispatch,
prompt_builder=_stub_prompt_builder,
log=harness_log,
run_id="run-factory",
concurrency=2,
max_steps=10,
skill_mode="auto",
plugins_factory=_factory,
)
assert len(factory_calls) == 2
call_set = set(factory_calls)
assert ("v1", "q1") in call_set
assert ("v2", "q2") in call_set
@pytest.mark.asyncio
async def test_no_factory_uses_empty_plugins(self, harness_log: HarnessLog) -> None:
"""plugins_factory=None 时使用空 plugins 列表。"""
llm = AsyncMock()
llm.chat.return_value = _make_llm_response(answer="B")
result = await run_inference(
[_make_question()],
llm=llm,
tool_dispatch_fn=_stub_tool_dispatch,
prompt_builder=_stub_prompt_builder,
log=harness_log,
run_id="run-no-factory",
concurrency=1,
max_steps=10,
skill_mode="auto",
plugins_factory=None,
)
assert result.total == 1
assert result.stop_reason_counts.get("finished") == 1
class TestConcurrencyControl:
"""并发控制 Semaphore 测试。"""
@pytest.mark.asyncio
async def test_concurrency_semaphore_limits(self, harness_log: HarnessLog) -> None:
"""Semaphore(1) 限制并发为 1 — 通过最大并发计数器验证。"""
import asyncio
llm = AsyncMock()
current_concurrent = 0
max_concurrent = 0
original_response = _make_llm_response(answer="B")
async def _slow_chat(
messages: Any,
*,
session_id: str | None = None,
parent_call_id: str | None = None,
) -> LLMResponse:
nonlocal current_concurrent, max_concurrent
current_concurrent += 1
max_concurrent = max(max_concurrent, current_concurrent)
await asyncio.sleep(0.01)
current_concurrent -= 1
return original_response
llm.chat.side_effect = _slow_chat
questions = [_make_question(question_id=f"q{i}") for i in range(5)]
await run_inference(
questions,
llm=llm,
tool_dispatch_fn=_stub_tool_dispatch,
prompt_builder=_stub_prompt_builder,
log=harness_log,
run_id="run-sem",
concurrency=1,
max_steps=10,
skill_mode="auto",
)
assert max_concurrent == 1
class TestTablesCreated:
"""表创建测试。"""
@pytest.mark.asyncio
async def test_five_tables_created(self, harness_log: HarnessLog) -> None:
"""run_inference 启动时创建 5 张推理表。"""
llm = AsyncMock()
await run_inference(
[],
llm=llm,
tool_dispatch_fn=_stub_tool_dispatch,
prompt_builder=_stub_prompt_builder,
log=harness_log,
run_id="run-tables",
concurrency=1,
max_steps=10,
skill_mode="auto",
)
expected_tables = [
"predictions",
"traces",
"validation_flags",
"anchor_check",
"observe_frame_health",
]
for table_name in expected_tables:
rows = harness_log.query(
"SELECT name FROM sqlite_master WHERE type='table' AND name=?",
(table_name,),
)
assert len(rows) == 1, f"表 {table_name} 未创建"