kb-skills/tests/conftest.py

"""共享测试 fixtures 和配置"""

import os
from typing import Optional
from unittest.mock import Mock
from pathlib import Path
from requests.exceptions import ConnectionError, Timeout, HTTPError

import pytest
from kb_workflow.srt.parser import Segment


class MockAIClient:
    """模拟 AI 客户端"""

    def __init__(self, response: str = "Mock AI response"):
        self.response = response

    def generate(
        self,
        prompt: str,
        max_tokens: int = 16000,
        disable_reasoning: Optional[bool] = None
    ) -> str:
        """生成模拟响应"""
        return self.response


@pytest.fixture
def mock_ai_client():
    """提供模拟 AI 客户端"""
    return MockAIClient()


@pytest.fixture
def sample_segments():
    """提供示例字幕分段数据"""
    return [
        {"index": 1, "start": "00:00:00,000", "end": "00:00:05,000", "text": "大家好，欢迎来到我的频道"},
        {"index": 2, "start": "00:00:05,000", "end": "00:00:10,000", "text": "今天我们要讲的是机器学习的基础知识"},
        {"index": 3, "start": "00:00:10,000", "end": "00:00:15,000", "text": "首先，什么是机器学习"},
        {"index": 4, "start": "00:00:15,000", "end": "00:00:20,000", "text": "机器学习是人工智能的一个分支"},
        {"index": 5, "start": "00:00:20,000", "end": "00:00:25,000", "text": "它让计算机能够从数据中学习"},
    ]


@pytest.fixture
def sample_metadata():
    """提供示例视频元数据"""
    return {
        "title": "机器学习入门教程",
        "category": "教育",
        "tags": ["编程", "AI", "教程"],
        "description": "从零开始学习机器学习"
    }


@pytest.fixture
def sample_metadata_movie():
    """提供电影类型示例元数据"""
    return {
        "title": "星际穿越",
        "category": "影视",
        "tags": ["科幻", "电影"],
        "description": "一部关于太空探索的科幻电影"
    }


# ============================================================================
# Phase 14 Fixtures: 鲁棒性和成本控制测试
# ============================================================================

@pytest.fixture
def long_text_sample():
    """长文本样本（>5000 字符，用于分块测试）"""
    # 生成一个超过 5000 字符的文本，包含多个句子
    sentences = ["这是一个测试句子。" * 10] * 15  # 约 6000 字符
    return " ".join(sentences)


@pytest.fixture
def mock_ai_client_network_error():
    """模拟网络错误的 AI 客户端"""
    client = Mock()
    client.api_key = "test_key"
    client.generate = Mock(side_effect=ConnectionError("Network error"))
    return client


@pytest.fixture
def mock_ai_client_timeout():
    """模拟超时的 AI 客户端"""
    client = Mock()
    client.api_key = "test_key"
    client.generate = Mock(side_effect=Timeout("Request timeout"))
    return client


@pytest.fixture
def mock_ai_client_rate_limit():
    """模拟限流的 AI 客户端"""
    response = Mock()
    response.status_code = 429
    client = Mock()
    client.api_key = "test_key"

    def generate_side_effect(*args, **kwargs):
        raise HTTPError(response=response)

    client.generate = Mock(side_effect=generate_side_effect)
    return client


@pytest.fixture
def mock_ai_client_api_error():
    """模拟 API 错误的 AI 客户端（不可恢复）"""
    response = Mock()
    response.status_code = 400
    client = Mock()
    client.api_key = "test_key"

    def generate_side_effect(*args, **kwargs):
        raise HTTPError(response=response)

    client.generate = Mock(side_effect=generate_side_effect)
    return client


@pytest.fixture
def mock_ai_client_parse_error():
    """模拟解析失败的 AI 客户端"""
    client = Mock()
    client.api_key = "test_key"
    # 第一次返回无效 JSON
    client.generate = Mock(return_value="Invalid JSON response")
    return client


@pytest.fixture
def mock_ai_client_success():
    """模拟成功的 AI 客户端"""
    import json
    client = Mock()
    client.api_key = "test_key"
    client.generate = Mock(return_value=json.dumps({
        "video_type": "知识类",
        "segments": [
            {
                "title": "测试标题",
                "content": "测试内容",
                "start_index": 0,
                "end_index": 10
            }
        ],
        "actual_segment_count": 1
    }, ensure_ascii=False))
    return client


@pytest.fixture
def temp_cache_dir(tmp_path):
    """临时缓存目录"""
    cache_dir = tmp_path / "cache" / "semantic"
    return cache_dir


@pytest.fixture
def semantic_cache(temp_cache_dir):
    """语义分段缓存实例"""
    from kb_workflow.srt.cache.semantic_cache import SemanticCache
    return SemanticCache(cache_dir=temp_cache_dir)


@pytest.fixture
def long_subtitle_segments():
    """长字幕分段样本（>5000 字符）"""
    segments = []
    for i in range(100):
        segments.append(Segment(
            id=i + 1,
            start_time=i * 1.0,
            end_time=(i + 1) * 1.0,
            text=f"这是第 {i + 1} 句字幕内容。" * 5
        ))
    return segments


# ============================================================================
# Phase 15 Fixtures: 配置集成测试
# ============================================================================

@pytest.fixture
def temp_env_vars(monkeypatch):
    """临时环境变量 fixture，测试后自动清理"""
    original_env = {}
    saved_vars = {}

    def _set_env(env_dict):
        nonlocal original_env, saved_vars
        # 保存原始值
        for key in env_dict.keys():
            original_env[key] = os.environ.get(key)
        # 设置新值
        for key, value in env_dict.items():
            if value is None:
                os.environ.pop(key, None)
            else:
                os.environ[key] = value
            saved_vars[key] = value

    yield _set_env

    # 清理：恢复原始环境变量
    for key, original_value in original_env.items():
        if original_value is None:
            os.environ.pop(key, None)
        else:
            os.environ[key] = original_value


@pytest.fixture
def mock_kbsettings():
    """Mock KBSettings fixture，用于隔离测试环境"""
    from kb_workflow.pipeline.config import KBSettings
    from kb_workflow.pipeline.config import get_settings

    mock_settings = MagicMock(spec=KBSettings)
    mock_settings.enable_semantic_segmentation = True
    mock_settings.base_path = "/tmp/kb"
    mock_settings.archive_path = "/tmp/archive"
    mock_settings.hf_token = MagicMock()
    mock_settings.hf_token.get_secret_value.return_value = ""
    mock_settings.zhihu_cookie = MagicMock()
    mock_settings.zhihu_cookie.get_secret_value.return_value = ""

    return mock_settings