diff --git a/tests/unit/test_build_trees.py b/tests/unit/test_build_trees.py new file mode 100644 index 0000000..7a03fda --- /dev/null +++ b/tests/unit/test_build_trees.py @@ -0,0 +1,108 @@ +"""tools/build_trees.py 单元测试。 + +覆盖纯函数（完整性校验、待建清单发现、SRT 查找）与编排集成（Task 3 追加）。 +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent +if str(PROJECT_ROOT) not in sys.path: + sys.path.insert(0, str(PROJECT_ROOT)) + +from app.tree.index import IndexMeta, L1Card, L1Node, TreeIndex # noqa: E402 +from tools.build_trees import ( # noqa: E402 + _discover_pending, + _find_srt_entries, + _tree_is_complete, +) + + +def _write_valid_tree(tree_path: Path) -> None: + """写入一棵最小合法树。""" + l1 = L1Node( + id="vid_L1_000", + card=L1Card("场景", "室内", ["实体"], ["动作"], ["关键词"], [], "线性"), + time_range=(0.0, 10.0), + children=[], + ) + index = TreeIndex(metadata=IndexMeta("/v.mp4", "video"), roots=[l1]) + tree_path.parent.mkdir(parents=True, exist_ok=True) + index.save_json(str(tree_path)) + + +class TestTreeIsComplete: + """_tree_is_complete 测试。""" + + def test_valid_tree(self, tmp_path: Path) -> None: + """合法 tree.json 判定完整。""" + tree_path = tmp_path / "vid" / "tree.json" + _write_valid_tree(tree_path) + assert _tree_is_complete(tree_path) is True + + def test_missing_file(self, tmp_path: Path) -> None: + """文件不存在判定不完整。""" + assert _tree_is_complete(tmp_path / "nope" / "tree.json") is False + + def test_corrupt_json(self, tmp_path: Path) -> None: + """损坏 JSON 判定不完整（不抛异常）。""" + p = tmp_path / "vid" / "tree.json" + p.parent.mkdir(parents=True) + p.write_text("{broken", encoding="utf-8") + assert _tree_is_complete(p) is False + + +class TestDiscoverPending: + """_discover_pending 测试。""" + + def _touch_videos(self, videos_dir: Path, names: list[str]) -> None: + videos_dir.mkdir(parents=True, exist_ok=True) + for n in names: + (videos_dir / n).write_bytes(b"") + + def test_all_pending_when_fresh(self, tmp_path: Path) -> None: + """无进度无产物时全部待建，按名排序。""" + videos = tmp_path / "videos" + self._touch_videos(videos, ["b.mp4", "a.mkv", "c.txt"]) + pending = _discover_pending(videos, tmp_path / "out", set()) + assert [p.name for p in pending] == ["a.mkv", "b.mp4"] # 非视频扩展名被忽略 + + def test_skips_finished_and_complete(self, tmp_path: Path) -> None: + """progress 已记录或 tree.json 完整的视频被跳过。""" + videos = tmp_path / "videos" + out = tmp_path / "out" + self._touch_videos(videos, ["a.mp4", "b.mp4", "c.mp4"]) + _write_valid_tree(out / "b" / "tree.json") # b 已有完整树 + pending = _discover_pending(videos, out, {"a"}) # a 在 progress 中 + assert [p.name for p in pending] == ["c.mp4"] + + def test_incomplete_tree_not_skipped(self, tmp_path: Path) -> None: + """tree.json 损坏的视频仍待建（重建覆盖）。""" + videos = tmp_path / "videos" + out = tmp_path / "out" + self._touch_videos(videos, ["a.mp4"]) + (out / "a").mkdir(parents=True) + (out / "a" / "tree.json").write_text("{broken", encoding="utf-8") + pending = _discover_pending(videos, out, set()) + assert [p.name for p in pending] == ["a.mp4"] + + +class TestFindSrtEntries: + """_find_srt_entries 测试。""" + + def test_found(self, tmp_path: Path) -> None: + """同名 .srt 存在时解析返回条目。""" + srt = tmp_path / "vid.srt" + srt.write_text( + "1\n00:00:01,000 --> 00:00:03,000\nhello world\n\n", + encoding="utf-8", + ) + entries = _find_srt_entries(tmp_path / "vid.mp4", tmp_path) + assert entries is not None + assert len(entries) == 1 + + def test_missing_returns_none(self, tmp_path: Path) -> None: + """无同名 .srt 返回 None。""" + assert _find_srt_entries(tmp_path / "vid.mp4", tmp_path) is None diff --git a/tools/build_trees.py b/tools/build_trees.py new file mode 100644 index 0000000..23fbe76 --- /dev/null +++ b/tools/build_trees.py @@ -0,0 +1,175 @@ +#!/usr/bin/env python3 +"""批量并行建树入口：多视频并发构建三层 TreeIndex。 + +并发模型（Spec-2）: + 视频级 Semaphore(video_concurrency) + gather —— 复刻 repair_trees.py 惯例； + 全局共享一个 API Semaphore(api_concurrency) 注入所有 VideoTreeBuilder， + 端点压力与单视频建树完全一致，吞吐提升来自非 API 阶段跨视频重叠。 + +用法: + conda activate Video-Tree-TRM + python tools/build_trees.py --videos-dir [--out-dir store/videos] + [--srt-dir ] [--video-concurrency 16] [--limit 0] + +api_concurrency 为工程配置，从 .env 读取 TREE_BUILD_API_CONCURRENCY（默认 16）。 +app/core/adapters 不 import 此脚本。 +""" + +from __future__ import annotations + +import json +import os +import sys +from pathlib import Path +from typing import TYPE_CHECKING + +# 确保项目根目录在 sys.path 中 +PROJECT_ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(PROJECT_ROOT)) + +from dotenv import load_dotenv # noqa: E402 +from loguru import logger # noqa: E402 + +load_dotenv(PROJECT_ROOT / ".env") + +from app.tree.index import TreeIndex # noqa: E402 +from app.tree.subtitle import parse_srt # noqa: E402 + +if TYPE_CHECKING: + import asyncio + + from app.tree.subtitle import SRTEntry + +# --------------------------------------------------------------------------- +# 日志配置：不缓存，立即输出 +# --------------------------------------------------------------------------- + +logger.remove() +logger.add( + sys.stderr, + format="{time:HH:mm:ss} | {level:<7} | {message}", + level="DEBUG", + colorize=True, +) +logger.add( + PROJECT_ROOT / "logs" / "build_trees.log", + format="{time:YYYY-MM-DD HH:mm:ss} | {level:<7} | {message}", + level="DEBUG", + rotation="50 MB", +) + +# --------------------------------------------------------------------------- +# 断点续跑 — progress 文件管理（复刻 repair_trees.py 惯例） +# --------------------------------------------------------------------------- + +PROGRESS_FILE = "build_progress.json" + +_VIDEO_SUFFIXES = frozenset({".mp4", ".mkv", ".avi", ".webm"}) + + +def load_progress(path: Path) -> set[str]: + """读取 progress 文件，返回已完成视频 ID 集合。 + + 参数: + path: progress JSON 文件路径。 + + 返回: + 已完成视频 ID 集合。文件不存在或损坏时返回空集。 + """ + if not path.exists(): + return set() + try: + data = json.loads(path.read_text(encoding="utf-8")) + return set(data.get("finished_video_ids", [])) + except (json.JSONDecodeError, KeyError, TypeError, AttributeError): + logger.warning("progress 文件损坏，忽略: {}", path) + return set() + + +async def save_progress(path: Path, lock: asyncio.Lock, vid: str) -> None: + """原子追加一个视频 ID 到 progress 文件。 + + 参数: + path: progress JSON 文件路径。 + lock: asyncio.Lock，防并发读改写丢更新。 + vid: 要追加的视频 ID。 + """ + async with lock: + finished = load_progress(path) + finished.add(vid) + tmp = path.with_suffix(".tmp") + tmp.write_text( + json.dumps({"finished_video_ids": sorted(finished)}, ensure_ascii=False, indent=2), + encoding="utf-8", + ) + os.replace(str(tmp), str(path)) + + +# --------------------------------------------------------------------------- +# 待建发现与完整性校验 +# --------------------------------------------------------------------------- + + +def _tree_is_complete(tree_path: Path) -> bool: + """判断 tree.json 是否存在且可加载为非空树。 + + 参数: + tree_path: tree.json 路径。 + + 返回: + True 表示完整（跳过重建）；文件缺失/损坏/空树返回 False。 + """ + if not tree_path.exists(): + return False + try: + index = TreeIndex.load_json(str(tree_path)) + except (json.JSONDecodeError, KeyError, TypeError, ValueError, AssertionError) as exc: + logger.warning("tree.json 无法加载，视为不完整: {} ({})", tree_path, exc) + return False + return len(index.roots) > 0 + + +def _discover_pending( + videos_dir: Path, + out_dir: Path, + finished: set[str], +) -> list[Path]: + """扫描视频目录，返回待建视频文件列表（按文件名排序）。 + + 跳过条件：video_id 在 progress 中，或 out_dir//tree.json 完整。 + + 参数: + videos_dir: 视频文件目录。 + out_dir: 树输出根目录。 + finished: progress 中已完成的视频 ID 集合。 + + 返回: + 待建视频文件路径列表。 + """ + pending: list[Path] = [] + for f in sorted(videos_dir.iterdir()): + if not f.is_file() or f.suffix.lower() not in _VIDEO_SUFFIXES: + continue + vid = f.stem + if vid in finished: + continue + if _tree_is_complete(out_dir / vid / "tree.json"): + continue + pending.append(f) + return pending + + +def _find_srt_entries(video_path: Path, srt_dir: Path) -> list[SRTEntry] | None: + """按视频同名规则查找并解析 SRT 字幕。 + + 参数: + video_path: 视频文件路径。 + srt_dir: SRT 目录。 + + 返回: + SRTEntry 列表；无同名 .srt 时返回 None。 + """ + srt_path = srt_dir / f"{video_path.stem}.srt" + if not srt_path.exists(): + return None + return parse_srt(str(srt_path))