#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ 测试WebSocket修复效果 """ import os import sys import time import requests import socketio # 设置环境变量 os.environ['PYTHONIOENCODING'] = 'utf-8' def test_websocket_connection(): """测试WebSocket连接""" print("=" * 60) print("🧪 WebSocket连接修复测试") print("=" * 60) # API服务器地址 api_base = "http://localhost:5000" # 测试HTTP连接 try: response = requests.get(f"{api_base}/api/health", timeout=5) if response.status_code == 200: print("✅ HTTP API服务器连接正常") else: print(f"⚠️ HTTP API服务器响应异常: {response.status_code}") return except Exception as e: print(f"❌ 无法连接到API服务器: {e}") print("💡 请先启动API服务器: PYTHONIOENCODING=utf-8 uv run server/api.py") return # 创建SocketIO客户端 sio = socketio.Client(logger=True, engineio_logger=False) # 连接状态 connected = False connection_error = None @sio.event def connect(): nonlocal connected connected = True print("✅ WebSocket连接成功") @sio.event def disconnect(): print("🔌 WebSocket连接断开") @sio.event def connection_established(data): print(f"📡 收到连接确认: {data}") @sio.event def training_progress_detailed(data): task_id = data.get('task_id', 'unknown')[:8] message = data.get('message', 'No message') print(f"📊 [训练进度] {task_id}: {message}") @sio.event def training_update(data): task_id = data.get('task_id', 'unknown')[:8] status = data.get('status', 'unknown') progress = data.get('progress', 0) print(f"🔄 [训练更新] {task_id}: {status} ({progress}%)") @sio.on('*') def catch_all(event, data): print(f"📨 收到事件: {event} -> {data}") # 尝试连接WebSocket try: print(f"\n🔗 尝试连接WebSocket: {api_base}") sio.connect(f"{api_base}", namespaces=['/training']) # 等待连接建立 timeout = 10 for i in range(timeout): if connected: break time.sleep(1) print(f"⏳ 等待连接... ({i+1}/{timeout})") if connected: print(f"🎉 WebSocket连接测试成功!") # 测试保持连接几秒钟 print(f"\n⏱️ 保持连接5秒钟,监听事件...") time.sleep(5) # 断开连接 sio.disconnect() print(f"👋 主动断开WebSocket连接") else: print(f"❌ WebSocket连接超时") except Exception as e: connection_error = str(e) print(f"❌ WebSocket连接失败: {e}") # 测试HTTP轮询备用方案 try: print(f"\n🔄 测试HTTP轮询方式...") sio_polling = socketio.Client(logger=False, engineio_logger=False) polling_connected = False @sio_polling.event def connect(): nonlocal polling_connected polling_connected = True print("✅ HTTP轮询连接成功") # 使用轮询传输方式 sio_polling.connect(f"{api_base}", transports=['polling'], namespaces=['/training']) # 等待连接 for i in range(5): if polling_connected: break time.sleep(1) if polling_connected: print("🎉 HTTP轮询备用方案工作正常") sio_polling.disconnect() else: print("⚠️ HTTP轮询也无法连接") except Exception as e: print(f"⚠️ HTTP轮询测试失败: {e}") print("\n" + "=" * 60) print("🎯 WebSocket测试结果:") if connected: print("✅ WebSocket协议: 正常") else: print("❌ WebSocket协议: 失败") if connection_error: print(f" 错误详情: {connection_error}") if polling_connected: print("✅ HTTP轮询备用: 正常") else: print("⚠️ HTTP轮询备用: 需要检查") print("\n💡 修复建议:") if not connected and not polling_connected: print("1. 检查Flask-SocketIO版本兼容性") print("2. 确认防火墙和端口设置") print("3. 检查服务器日志输出") elif connected: print("✅ WebSocket修复成功,可以正常使用实时功能") print("=" * 60) if __name__ == "__main__": test_websocket_connection()