TermControlAgent/test_api_only.py

#!/usr/bin/env python3
"""
专门测试API端点的脚本
"""
import asyncio
import json
import sys
import os

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

from app.core.websocket.manager import websocket_manager
from app.schemas.websocket import CreateWebSocketClientRequest
from app.utils.log import get_logger

logger = get_logger(__name__)

async def test_api_endpoint_simulation():
    """模拟API端点的完整流程"""
    print("=== API端点模拟测试 ===")
    
    # 模拟API请求
    test_request = {
        "name": "test_1",
        "url": "ws://localhost:8080/ws",
        "heartbeat_interval": 120
    }
    
    print(f"模拟API请求: {json.dumps(test_request, indent=2)}")
    
    try:
        # 1. 创建请求对象（模拟API验证）
        request = CreateWebSocketClientRequest(**test_request)
        print("✅ 请求验证通过")
        
        # 2. 创建客户端
        print("1. 创建WebSocket客户端...")
        client = await websocket_manager.create_client(request.name, request.url)
        print(f"   ✅ 客户端创建成功: {client.name}")
        
        # 3. 连接客户端
        print("2. 连接WebSocket客户端...")
        success = await websocket_manager.connect_client(request.name)
        
        if success:
            print(f"   ✅ 连接成功: {request.name}")
            
            # 4. 确保默认Channels存在并连接，创建适配器
            default_channels = ["default", "system", "events"]
            max_channel_size = 1000
            
            print("3. 创建默认Channels和适配器...")
            for ch in default_channels:
                channel = websocket_manager.get_channel(ch)
                if not channel:
                    channel = await websocket_manager.create_channel(ch, max_channel_size)
                await channel.connect()
                await websocket_manager.create_adapter(request.name, ch, request.heartbeat_interval)
                print(f"   ✅ Channel {ch} 和适配器创建成功")
            
            # 5. 检查最终状态
            print("4. 检查最终状态...")
            client = websocket_manager.get_client(request.name)
            if client:
                stats = client.get_stats()
                print(f"   客户端状态: {stats['state']}")
                print(f"   连接状态: {stats['is_connected']}")
                print(f"   重连次数: {stats['reconnect_attempts']}")
            
            # 6. 检查适配器状态
            print("5. 检查适配器状态...")
            for ch in default_channels:
                adapter = websocket_manager.get_adapter(request.name, ch)
                if adapter:
                    adapter_stats = adapter.get_stats()
                    print(f"   适配器 {ch}: 发送任务运行={adapter_stats['send_task_running']}")
            
            # 7. 模拟成功响应
            response_data = {
                "name": request.name, 
                "url": request.url, 
                "status": "connected", 
                "heartbeat_interval": request.heartbeat_interval
            }
            print(f"✅ API响应: {json.dumps(response_data, indent=2)}")
            
            # 8. 等待一段时间观察稳定性
            print("6. 等待5秒观察连接稳定性...")
            await asyncio.sleep(5)
            
            # 9. 再次检查状态
            client = websocket_manager.get_client(request.name)
            if client:
                stats = client.get_stats()
                print(f"   5秒后状态: {stats['state']}")
                print(f"   5秒后连接: {stats['is_connected']}")
            
            # 10. 断开连接
            print("7. 断开WebSocket客户端...")
            await websocket_manager.disconnect_client(request.name)
            print(f"   ✅ 断开成功: {request.name}")
            
        else:
            print(f"   ❌ 连接失败: {request.name}")
            # 模拟API错误响应
            error_response = {
                "success": False,
                "message": f"WebSocket客户端 {request.name} 创建或连接失败",
                "error_code": "CONNECTION_FAILED"
            }
            print(f"❌ API错误响应: {json.dumps(error_response, indent=2)}")
            
    except Exception as e:
        print(f"   ❌ API测试异常: {e}")
        logger.error(f"API测试异常: {e}")
        # 模拟API异常响应
        error_response = {
            "success": False,
            "message": str(e),
            "error_code": "INTERNAL_ERROR"
        }
        print(f"❌ API异常响应: {json.dumps(error_response, indent=2)}")
    
    # 清理
    try:
        await websocket_manager.remove_client("test_1")
    except:
        pass

async def main():
    """主函数"""
    try:
        await test_api_endpoint_simulation()
        
    except KeyboardInterrupt:
        print("\n测试被用户中断")
    except Exception as e:
        print(f"测试过程中发生异常: {e}")
        logger.error(f"测试异常: {e}")
    finally:
        # 清理资源
        print("\n清理WebSocket管理器资源...")
        await websocket_manager.cleanup()
        print("清理完成")

if __name__ == "__main__":
    # 运行测试
    asyncio.run(main())