kb-skills/tests/test_semantic_engine.py

"""测试语义分段引擎"""

import pytest

from kb_workflow.srt.semantic_engine import (
    VideoSegment,
    SegmentationResult,
    SemanticSegmentationEngine
)


class TestVideoTypeDetection:
    """测试视频类型识别功能"""

    def test_video_type_detection_from_title(self):
        """测试从标题推断视频类型"""
        # 知识类 - 从标题推断
        metadata_edu = {"title": "Python编程教程", "category": "教育"}
        engine = SemanticSegmentationEngine()
        video_type = engine._infer_video_type(metadata_edu)
        assert video_type == "知识类"

        # 知识类 - 包含"课程"
        metadata_course = {"title": "机器学习课程", "category": ""}
        video_type = engine._infer_video_type(metadata_course)
        assert video_type == "知识类"

    def test_video_type_detection_from_category(self):
        """测试从分类推断视频类型"""
        # 故事类 - 从分类推断
        metadata_movie = {"title": "星际穿越", "category": "影视"}
        engine = SemanticSegmentationEngine()
        video_type = engine._infer_video_type(metadata_movie)
        assert video_type == "故事类"

        # 访谈类 - 从标签推断
        metadata_interview = {"title": "名人访谈", "tags": ["访谈", "对话"]}
        video_type = engine._infer_video_type(metadata_interview)
        assert video_type == "访谈类"

    def test_video_type_fallback_to_default(self):
        """测试元数据缺失时使用默认类型"""
        # 空元数据
        engine = SemanticSegmentationEngine()
        video_type = engine._infer_video_type({})
        assert video_type == "知识类"  # 默认类型

        # 无效元数据
        video_type = engine._infer_video_type({"title": "", "category": ""})
        assert video_type == "知识类"


class TestVideoTypeInPrompt:
    """测试视频类型在 Prompt 中的使用"""

    def test_video_type_in_prompt(self):
        """测试 Prompt 包含视频类型信息"""
        engine = SemanticSegmentationEngine()
        metadata = {"title": "机器学习教程", "category": "教育"}

        prompt = engine._build_prompt("测试内容", metadata)

        assert "视频类型" in prompt or "知识类" in prompt
        assert "机器学习教程" in prompt

    def test_prompt_format(self):
        """测试 Prompt 格式正确"""
        engine = SemanticSegmentationEngine()
        content = "这是测试内容"
        metadata = {"title": "测试视频"}

        prompt = engine._build_prompt(content, metadata)

        assert "请分析以下视频字幕" in prompt or "字幕内容" in prompt
        assert content in prompt


class TestPydanticModels:
    """测试 Pydantic 模型验证"""

    def test_video_segment_model(self):
        """测试 VideoSegment 模型"""
        # 创建完整分段
        segment = VideoSegment(
            title="机器学习概述",
            content="机器学习是人工智能的一个分支",
            start_index=1,
            end_index=5
        )
        assert segment.title == "机器学习概述"
        assert segment.content == "机器学习是人工智能的一个分支"
        assert segment.start_index == 1
        assert segment.end_index == 5

        # 创建无标题分段（title 可选）
        segment_no_title = VideoSegment(
            content="测试内容",
            start_index=1,
            end_index=3
        )
        assert segment_no_title.title is None

    def test_segmentation_result_model(self):
        """测试 SegmentationResult 模型"""
        segments = [
            VideoSegment(title="第一章", content="内容1", start_index=1, end_index=5),
            VideoSegment(title="第二章", content="内容2", start_index=6, end_index=10)
        ]

        result = SegmentationResult(
            video_type="知识类",
            segments=segments,
            actual_segment_count=2
        )

        assert result.video_type == "知识类"
        assert len(result.segments) == 2
        assert result.actual_segment_count == 2

    def test_model_validation(self):
        """测试模型字段验证"""
        # 缺少必填字段应该抛出异常
        with pytest.raises(Exception):  # Pydantic ValidationError
            VideoSegment(
                title="测试",
                # 缺少 content
                start_index=1,
                end_index=3
            )