LocalAgent/tests/test_retry_fix.py

"""
测试重试策略修复
验证网络异常能够被正确识别并重试
"""

import sys
import io
from pathlib import Path

# 设置标准输出为 UTF-8
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')

# 添加项目根目录到路径
PROJECT_ROOT = Path(__file__).parent
sys.path.insert(0, str(PROJECT_ROOT))

from llm.client import LLMClient, LLMClientError
import requests


def test_exception_classification():
    """测试异常分类"""
    print("=" * 60)
    print("测试 1: 异常分类")
    print("=" * 60)
    
    # 测试网络异常
    network_error = LLMClientError(
        "网络连接失败",
        error_type=LLMClientError.TYPE_NETWORK,
        original_exception=requests.exceptions.ConnectionError()
    )
    print(f"✓ 网络错误类型: {network_error.error_type}")
    assert network_error.error_type == LLMClientError.TYPE_NETWORK
    
    # 测试服务器异常
    server_error = LLMClientError(
        "服务器错误 500",
        error_type=LLMClientError.TYPE_SERVER
    )
    print(f"✓ 服务器错误类型: {server_error.error_type}")
    assert server_error.error_type == LLMClientError.TYPE_SERVER
    
    # 测试客户端异常
    client_error = LLMClientError(
        "请求参数错误 400",
        error_type=LLMClientError.TYPE_CLIENT
    )
    print(f"✓ 客户端错误类型: {client_error.error_type}")
    assert client_error.error_type == LLMClientError.TYPE_CLIENT
    
    print("\n✅ 异常分类测试通过\n")


def test_should_retry_logic():
    """测试重试判断逻辑"""
    print("=" * 60)
    print("测试 2: 重试判断逻辑")
    print("=" * 60)
    
    client = LLMClient(max_retries=3)
    
    # 测试网络错误应该重试
    network_error = LLMClientError(
        "网络连接失败",
        error_type=LLMClientError.TYPE_NETWORK,
        original_exception=requests.exceptions.ConnectionError()
    )
    should_retry = client._should_retry(network_error)
    print(f"✓ 网络错误应该重试: {should_retry}")
    assert should_retry == True, "网络错误应该重试"
    
    # 测试超时错误应该重试
    timeout_error = LLMClientError(
        "请求超时",
        error_type=LLMClientError.TYPE_NETWORK,
        original_exception=requests.exceptions.Timeout()
    )
    should_retry = client._should_retry(timeout_error)
    print(f"✓ 超时错误应该重试: {should_retry}")
    assert should_retry == True, "超时错误应该重试"
    
    # 测试服务器错误应该重试
    server_error = LLMClientError(
        "服务器错误 500",
        error_type=LLMClientError.TYPE_SERVER
    )
    should_retry = client._should_retry(server_error)
    print(f"✓ 服务器错误应该重试: {should_retry}")
    assert should_retry == True, "服务器错误应该重试"
    
    # 测试客户端错误不应该重试
    client_error = LLMClientError(
        "请求参数错误 400",
        error_type=LLMClientError.TYPE_CLIENT
    )
    should_retry = client._should_retry(client_error)
    print(f"✓ 客户端错误不应该重试: {should_retry}")
    assert should_retry == False, "客户端错误不应该重试"
    
    # 测试解析错误不应该重试
    parse_error = LLMClientError(
        "解析响应失败",
        error_type=LLMClientError.TYPE_PARSE
    )
    should_retry = client._should_retry(parse_error)
    print(f"✓ 解析错误不应该重试: {should_retry}")
    assert should_retry == False, "解析错误不应该重试"
    
    # 测试配置错误不应该重试
    config_error = LLMClientError(
        "未配置 API Key",
        error_type=LLMClientError.TYPE_CONFIG
    )
    should_retry = client._should_retry(config_error)
    print(f"✓ 配置错误不应该重试: {should_retry}")
    assert should_retry == False, "配置错误不应该重试"
    
    # 测试原始异常检查
    error_with_original = LLMClientError(
        "网络请求异常",
        error_type=LLMClientError.TYPE_NETWORK,
        original_exception=requests.exceptions.ConnectionError("Connection refused")
    )
    should_retry = client._should_retry(error_with_original)
    print(f"✓ 带原始异常的网络错误应该重试: {should_retry}")
    assert should_retry == True, "带原始异常的网络错误应该重试"
    
    print("\n✅ 重试判断逻辑测试通过\n")


def test_error_type_preservation():
    """测试错误类型保留"""
    print("=" * 60)
    print("测试 3: 错误类型保留")
    print("=" * 60)
    
    # 模拟不同状态码的错误
    test_cases = [
        (500, LLMClientError.TYPE_SERVER, "服务器错误"),
        (502, LLMClientError.TYPE_SERVER, "网关错误"),
        (503, LLMClientError.TYPE_SERVER, "服务不可用"),
        (504, LLMClientError.TYPE_SERVER, "网关超时"),
        (429, LLMClientError.TYPE_SERVER, "限流错误"),
        (400, LLMClientError.TYPE_CLIENT, "请求错误"),
        (401, LLMClientError.TYPE_CLIENT, "未授权"),
        (403, LLMClientError.TYPE_CLIENT, "禁止访问"),
        (404, LLMClientError.TYPE_CLIENT, "未找到"),
    ]
    
    for status_code, expected_type, description in test_cases:
        if status_code >= 500:
            error_type = LLMClientError.TYPE_SERVER
        elif status_code == 429:
            error_type = LLMClientError.TYPE_SERVER
        else:
            error_type = LLMClientError.TYPE_CLIENT
        
        print(f"✓ 状态码 {status_code} ({description}): {error_type}")
        assert error_type == expected_type, f"状态码 {status_code} 的错误类型不正确"
    
    print("\n✅ 错误类型保留测试通过\n")


def main():
    """运行所有测试"""
    print("\n" + "=" * 60)
    print("重试策略修复验证测试")
    print("=" * 60 + "\n")
    
    try:
        test_exception_classification()
        test_should_retry_logic()
        test_error_type_preservation()
        
        print("=" * 60)
        print("✅ 所有测试通过！")
        print("=" * 60)
        print("\n修复总结：")
        print("1. ✅ 为 LLMClientError 添加了错误类型分类")
        print("2. ✅ 保留了原始异常信息")
        print("3. ✅ 统一了 _should_retry 判断逻辑")
        print("4. ✅ 网络异常（超时、连接失败）现在可以正确重试")
        print("5. ✅ 服务器错误（5xx）和限流（429）可以重试")
        print("6. ✅ 客户端错误（4xx）、解析错误、配置错误不会重试")
        print("7. ✅ 增强了重试度量指标记录")
        print("\n预期效果：")
        print("- 弱网环境下的稳定性显著提升")
        print("- 意图识别、生成计划、代码生成的成功率提高")
        print("- 网络抖动时自动重试并恢复")
        
    except AssertionError as e:
        print(f"\n❌ 测试失败: {e}")
        sys.exit(1)
    except Exception as e:
        print(f"\n❌ 测试出错: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()