iomgaa
6bdb802f01
- Add all claude skills (brainstorming, commit, debugging, TDD, etc.) - Add claude hooks (pre-commit-guard, post-edit-quality) - Add research templates (experiment plan, research brief, etc.) - Add claude tools (arxiv/semantic_scholar/openalex fetch, wiki, exa) - Add TRM4 reference implementation as algorithm fidelity baseline - Add research-wiki content (plans, index, graph, query_pack) - Update .gitignore to exclude .graphify_version runtime state
358 lines
12 KiB
Python
358 lines
12 KiB
Python
"""
|
||
test_pipeline.py — Pipeline 单元测试
|
||
======================================
|
||
使用 unittest.mock.MagicMock + patch 隔离所有外部依赖（无真实 API / 文件 IO）。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
from pathlib import Path
|
||
from unittest.mock import MagicMock, patch
|
||
|
||
import numpy as np
|
||
import pytest
|
||
import torch
|
||
|
||
from video_tree_trm.pipeline import Pipeline
|
||
from video_tree_trm.tree_index import IndexMeta, L1Node, L2Node, L3Node, TreeIndex
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 辅助：构造最小 Config Mock
|
||
# ---------------------------------------------------------------------------
|
||
|
||
D = 8 # 嵌入维度（与 RetrieverConfig 一致）
|
||
|
||
|
||
def _make_config(checkpoint: str | None = None) -> MagicMock:
|
||
"""返回一个 Mock Config，字段值与实际 dataclass 对齐。"""
|
||
cfg = MagicMock()
|
||
cfg.embed.model_name = "test-embed"
|
||
cfg.embed.embed_dim = D
|
||
cfg.retriever.checkpoint = checkpoint
|
||
cfg.retriever.embed_dim = D
|
||
cfg.retriever.num_heads = 2
|
||
cfg.retriever.L_layers = 2
|
||
cfg.retriever.L_cycles = 2
|
||
cfg.retriever.max_rounds = 2
|
||
cfg.retriever.ffn_expansion = 2.0
|
||
cfg.tree.cache_dir = "/tmp/test_pipeline_cache"
|
||
return cfg
|
||
|
||
|
||
def _make_small_tree() -> TreeIndex:
|
||
"""构造最小 1×1×1 TreeIndex，用于 query() 测试。"""
|
||
meta = IndexMeta(
|
||
source_path="dummy",
|
||
modality="text",
|
||
embed_model="test",
|
||
embed_dim=D,
|
||
)
|
||
l3 = L3Node(
|
||
id="l3_0",
|
||
description="节点描述",
|
||
embedding=np.zeros(D, dtype=np.float32),
|
||
raw_content="节点内容",
|
||
)
|
||
l2 = L2Node(
|
||
id="l2_0",
|
||
description="L2",
|
||
embedding=np.zeros(D, dtype=np.float32),
|
||
children=[l3],
|
||
)
|
||
l1 = L1Node(
|
||
id="l1_0",
|
||
summary="L1",
|
||
embedding=np.zeros(D, dtype=np.float32),
|
||
children=[l2],
|
||
)
|
||
return TreeIndex(metadata=meta, roots=[l1])
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Patch 工厂：将所有子模块构造函数替换为 MagicMock
|
||
# ---------------------------------------------------------------------------
|
||
|
||
_PATCHES = [
|
||
"video_tree_trm.pipeline.EmbeddingModel",
|
||
"video_tree_trm.pipeline.LLMClient",
|
||
"video_tree_trm.pipeline.RecursiveRetriever",
|
||
"video_tree_trm.pipeline.AnswerGenerator",
|
||
"video_tree_trm.pipeline.TextTreeBuilder",
|
||
"video_tree_trm.pipeline.VideoTreeBuilder",
|
||
]
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Pipeline.__init__ 测试
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def test_pipeline_init_components() -> None:
|
||
"""__init__ 后各属性（embed_model/llm/vlm/retriever/generator）均存在。"""
|
||
cfg = _make_config()
|
||
with patch.multiple(
|
||
"video_tree_trm.pipeline",
|
||
EmbeddingModel=MagicMock(),
|
||
LLMClient=MagicMock(),
|
||
RecursiveRetriever=MagicMock(),
|
||
AnswerGenerator=MagicMock(),
|
||
):
|
||
p = Pipeline(cfg)
|
||
|
||
assert hasattr(p, "embed_model"), "缺少 embed_model 属性"
|
||
assert hasattr(p, "llm"), "缺少 llm 属性"
|
||
assert hasattr(p, "vlm"), "缺少 vlm 属性"
|
||
assert hasattr(p, "retriever"), "缺少 retriever 属性"
|
||
assert hasattr(p, "generator"), "缺少 generator 属性"
|
||
|
||
|
||
def test_pipeline_init_no_checkpoint() -> None:
|
||
"""checkpoint=None 时 load_state_dict 不被调用。"""
|
||
cfg = _make_config(checkpoint=None)
|
||
mock_retriever_instance = MagicMock()
|
||
MockRetriever = MagicMock(return_value=mock_retriever_instance)
|
||
|
||
with patch.multiple(
|
||
"video_tree_trm.pipeline",
|
||
EmbeddingModel=MagicMock(),
|
||
LLMClient=MagicMock(),
|
||
RecursiveRetriever=MockRetriever,
|
||
AnswerGenerator=MagicMock(),
|
||
):
|
||
Pipeline(cfg)
|
||
|
||
mock_retriever_instance.load_state_dict.assert_not_called()
|
||
|
||
|
||
def test_pipeline_init_with_checkpoint(tmp_path: Path) -> None:
|
||
"""checkpoint 非 None 时 load_state_dict 被调用一次。"""
|
||
ckpt_file = tmp_path / "model.pt"
|
||
ckpt_file.write_bytes(b"") # 创建空文件，使 os.path.isfile 返回 True
|
||
|
||
cfg = _make_config(checkpoint=str(ckpt_file))
|
||
mock_retriever_instance = MagicMock()
|
||
MockRetriever = MagicMock(return_value=mock_retriever_instance)
|
||
|
||
fake_state_dict = {"weight": torch.zeros(1)}
|
||
|
||
with patch.multiple(
|
||
"video_tree_trm.pipeline",
|
||
EmbeddingModel=MagicMock(),
|
||
LLMClient=MagicMock(),
|
||
RecursiveRetriever=MockRetriever,
|
||
AnswerGenerator=MagicMock(),
|
||
), patch("video_tree_trm.pipeline.torch.load", return_value=fake_state_dict):
|
||
Pipeline(cfg)
|
||
|
||
mock_retriever_instance.load_state_dict.assert_called_once_with(fake_state_dict)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Pipeline.build_index 测试
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def test_build_index_text_calls_builder(tmp_path: Path) -> None:
|
||
"""文本模式调用 TextTreeBuilder.build，参数含文件内容。"""
|
||
src = tmp_path / "doc.txt"
|
||
src.write_text("文档内容", encoding="utf-8")
|
||
|
||
cfg = _make_config()
|
||
cfg.tree.cache_dir = str(tmp_path / "cache")
|
||
|
||
mock_tree = MagicMock(spec=TreeIndex)
|
||
mock_builder_instance = MagicMock()
|
||
mock_builder_instance.build.return_value = mock_tree
|
||
MockTextBuilder = MagicMock(return_value=mock_builder_instance)
|
||
|
||
with patch.multiple(
|
||
"video_tree_trm.pipeline",
|
||
EmbeddingModel=MagicMock(),
|
||
LLMClient=MagicMock(),
|
||
RecursiveRetriever=MagicMock(),
|
||
AnswerGenerator=MagicMock(),
|
||
TextTreeBuilder=MockTextBuilder,
|
||
):
|
||
p = Pipeline(cfg)
|
||
result = p.build_index(str(src), modality="text")
|
||
|
||
mock_builder_instance.build.assert_called_once()
|
||
call_args = mock_builder_instance.build.call_args
|
||
assert "文档内容" in call_args[0][0], "TextTreeBuilder.build 应传入文件内容"
|
||
assert result is mock_tree
|
||
|
||
|
||
def test_build_index_video_calls_builder(tmp_path: Path) -> None:
|
||
"""视频模式调用 VideoTreeBuilder.build，参数为 source_path。"""
|
||
cfg = _make_config()
|
||
cfg.tree.cache_dir = str(tmp_path / "cache")
|
||
|
||
mock_tree = MagicMock(spec=TreeIndex)
|
||
mock_builder_instance = MagicMock()
|
||
mock_builder_instance.build.return_value = mock_tree
|
||
MockVideoBuilder = MagicMock(return_value=mock_builder_instance)
|
||
|
||
video_path = "/fake/video.mp4"
|
||
|
||
with patch.multiple(
|
||
"video_tree_trm.pipeline",
|
||
EmbeddingModel=MagicMock(),
|
||
LLMClient=MagicMock(),
|
||
RecursiveRetriever=MagicMock(),
|
||
AnswerGenerator=MagicMock(),
|
||
VideoTreeBuilder=MockVideoBuilder,
|
||
):
|
||
p = Pipeline(cfg)
|
||
result = p.build_index(video_path, modality="video")
|
||
|
||
mock_builder_instance.build.assert_called_once_with(video_path)
|
||
assert result is mock_tree
|
||
|
||
|
||
def test_build_index_cache_hit(tmp_path: Path) -> None:
|
||
"""缓存文件存在时直接 TreeIndex.load，不重新构建。"""
|
||
cfg = _make_config()
|
||
cache_dir = tmp_path / "cache"
|
||
cache_dir.mkdir()
|
||
cfg.tree.cache_dir = str(cache_dir)
|
||
|
||
# 手动创建缓存文件（空文件即可让 isfile 返回 True）
|
||
cache_file = cache_dir / "doc_text.pkl"
|
||
cache_file.write_bytes(b"")
|
||
|
||
mock_tree = MagicMock(spec=TreeIndex)
|
||
mock_text_builder = MagicMock()
|
||
|
||
with patch.multiple(
|
||
"video_tree_trm.pipeline",
|
||
EmbeddingModel=MagicMock(),
|
||
LLMClient=MagicMock(),
|
||
RecursiveRetriever=MagicMock(),
|
||
AnswerGenerator=MagicMock(),
|
||
TextTreeBuilder=mock_text_builder,
|
||
), patch("video_tree_trm.pipeline.TreeIndex.load", return_value=mock_tree) as mock_load:
|
||
p = Pipeline(cfg)
|
||
result = p.build_index(str(tmp_path / "doc.txt"), modality="text")
|
||
|
||
mock_load.assert_called_once_with(str(cache_file))
|
||
mock_text_builder.return_value.build.assert_not_called()
|
||
assert result is mock_tree
|
||
|
||
|
||
def test_build_index_saves_cache(tmp_path: Path) -> None:
|
||
"""缓存不存在时构建后调用 tree.save。"""
|
||
cfg = _make_config()
|
||
cfg.tree.cache_dir = str(tmp_path / "cache")
|
||
|
||
src = tmp_path / "doc.txt"
|
||
src.write_text("内容", encoding="utf-8")
|
||
|
||
mock_tree = MagicMock(spec=TreeIndex)
|
||
mock_builder_instance = MagicMock()
|
||
mock_builder_instance.build.return_value = mock_tree
|
||
|
||
with patch.multiple(
|
||
"video_tree_trm.pipeline",
|
||
EmbeddingModel=MagicMock(),
|
||
LLMClient=MagicMock(),
|
||
RecursiveRetriever=MagicMock(),
|
||
AnswerGenerator=MagicMock(),
|
||
TextTreeBuilder=MagicMock(return_value=mock_builder_instance),
|
||
):
|
||
p = Pipeline(cfg)
|
||
p.build_index(str(src), modality="text")
|
||
|
||
mock_tree.save.assert_called_once()
|
||
saved_path: str = mock_tree.save.call_args[0][0]
|
||
assert "doc_text.pkl" in saved_path, f"保存路径应含 'doc_text.pkl'，实际={saved_path}"
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Pipeline.query 测试
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def test_query_embeds_question() -> None:
|
||
"""query() 调用 embed_model.embed_tensor(question)。"""
|
||
cfg = _make_config()
|
||
tree = _make_small_tree()
|
||
|
||
mock_embed = MagicMock()
|
||
mock_embed.embed_tensor.return_value = torch.zeros(1, D)
|
||
MockEmbed = MagicMock(return_value=mock_embed)
|
||
|
||
mock_retriever_instance = MagicMock()
|
||
mock_retriever_instance.return_value = {"paths": [(0, 0, 0)], "num_rounds": 1}
|
||
|
||
with patch.multiple(
|
||
"video_tree_trm.pipeline",
|
||
EmbeddingModel=MockEmbed,
|
||
LLMClient=MagicMock(),
|
||
RecursiveRetriever=MagicMock(return_value=mock_retriever_instance),
|
||
AnswerGenerator=MagicMock(),
|
||
):
|
||
p = Pipeline(cfg)
|
||
p.query("测试问题", tree)
|
||
|
||
mock_embed.embed_tensor.assert_called_once_with("测试问题")
|
||
|
||
|
||
def test_query_calls_retriever() -> None:
|
||
"""query() 调用 retriever(q, tree)。"""
|
||
cfg = _make_config()
|
||
tree = _make_small_tree()
|
||
|
||
q_tensor = torch.zeros(1, D)
|
||
mock_embed = MagicMock()
|
||
mock_embed.embed_tensor.return_value = q_tensor
|
||
|
||
mock_retriever_instance = MagicMock()
|
||
mock_retriever_instance.return_value = {"paths": [(0, 0, 0)], "num_rounds": 1}
|
||
|
||
with patch.multiple(
|
||
"video_tree_trm.pipeline",
|
||
EmbeddingModel=MagicMock(return_value=mock_embed),
|
||
LLMClient=MagicMock(),
|
||
RecursiveRetriever=MagicMock(return_value=mock_retriever_instance),
|
||
AnswerGenerator=MagicMock(),
|
||
):
|
||
p = Pipeline(cfg)
|
||
p.query("测试问题", tree)
|
||
|
||
mock_retriever_instance.assert_called_once()
|
||
call_args = mock_retriever_instance.call_args
|
||
# 第一个位置参数应为嵌入 Tensor，第二个为 tree
|
||
assert call_args[0][1] is tree, "retriever 第二个参数应为 tree"
|
||
|
||
|
||
def test_query_returns_answer() -> None:
|
||
"""query() 返回 generator.generate 的返回值。"""
|
||
cfg = _make_config()
|
||
tree = _make_small_tree()
|
||
|
||
mock_embed = MagicMock()
|
||
mock_embed.embed_tensor.return_value = torch.zeros(1, D)
|
||
|
||
mock_retriever_instance = MagicMock()
|
||
mock_retriever_instance.return_value = {"paths": [(0, 0, 0)], "num_rounds": 1}
|
||
|
||
mock_generator_instance = MagicMock()
|
||
mock_generator_instance.generate.return_value = "生成的答案"
|
||
|
||
with patch.multiple(
|
||
"video_tree_trm.pipeline",
|
||
EmbeddingModel=MagicMock(return_value=mock_embed),
|
||
LLMClient=MagicMock(),
|
||
RecursiveRetriever=MagicMock(return_value=mock_retriever_instance),
|
||
AnswerGenerator=MagicMock(return_value=mock_generator_instance),
|
||
):
|
||
p = Pipeline(cfg)
|
||
answer = p.query("问题", tree)
|
||
|
||
assert answer == "生成的答案", f"query() 应返回 generator 的结果，实际='{answer}'"
|
||
mock_generator_instance.generate.assert_called_once_with(
|
||
"问题", [(0, 0, 0)], tree
|
||
)
|