TermControlAgent/tests/test_websocket_ssl.py

#!/usr/bin/env python3
"""
WebSocket SSL证书验证测试脚本
"""
import asyncio
import sys
import os

# 添加项目根目录到Python路径
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from app.core.websocket.client import WebSocketClient
from app.core.config.settings import config
from app.utils.structured_log import get_structured_logger, LogLevel

logger = get_structured_logger(__name__, LogLevel.DEBUG)

async def test_websocket_ssl():
    """测试WebSocket SSL连接"""
    print("=== WebSocket SSL证书验证测试 ===")
    
    # 显示当前配置
    print(f"SSL证书验证: {config.websocket.ssl_verify_certificate}")
    print(f"SSL主机名验证: {config.websocket.ssl_verify_hostname}")
    
    # 测试URL（使用自签名证书的WebSocket服务器）
    test_url = "wss://echo.websocket.org"  # 使用公共测试服务器
    
    print(f"\n测试连接到: {test_url}")
    
    try:
        # 创建WebSocket客户端
        client = WebSocketClient(test_url, "test_client")
        
        # 尝试连接
        print("正在连接...")
        success = await client.connect()
        
        if success:
            print("✅ 连接成功！SSL证书验证问题已修复")
            
            # 发送测试消息
            test_message = {
                "type": "test",
                "data": "Hello WebSocket!"
            }
            
            # 注册消息处理器
            def handle_response(message):
                print(f"收到响应: {message}")
            
            client.register_message_handler("test", handle_response)
            
            # 等待一段时间接收响应
            await asyncio.sleep(2)
            
            # 断开连接
            await client.disconnect()
            print("连接已断开")
        else:
            print("❌ 连接失败")
            
    except Exception as e:
        print(f"❌ 测试过程中出现异常: {e}")
        logger.error(f"WebSocket SSL测试失败: {e}")

async def test_config_override():
    """测试配置覆盖功能"""
    print("\n=== 配置覆盖测试 ===")
    
    # 临时修改配置
    original_verify_cert = config.websocket.ssl_verify_certificate
    original_verify_host = config.websocket.ssl_verify_hostname
    
    try:
        # 设置为启用验证（应该会失败）
        config.websocket.ssl_verify_certificate = True
        config.websocket.ssl_verify_hostname = True
        
        print("配置已设置为启用SSL验证")
        print(f"SSL证书验证: {config.websocket.ssl_verify_certificate}")
        print(f"SSL主机名验证: {config.websocket.ssl_verify_hostname}")
        
        # 这里可以添加更多测试逻辑
        
    finally:
        # 恢复原始配置
        config.websocket.ssl_verify_certificate = original_verify_cert
        config.websocket.ssl_verify_hostname = original_verify_host
        print("配置已恢复")

if __name__ == "__main__":
    print("开始WebSocket SSL测试...")
    
    # 运行测试
    asyncio.run(test_websocket_ssl())
    asyncio.run(test_config_override())
    
    print("\n测试完成！")