TermControlAgent/test_heartbeat_simple.py

#!/usr/bin/env python3
"""
简单的心跳功能测试
验证心跳任务是否正确启动并发送心跳消息
"""
import asyncio
import sys
import os
from datetime import datetime

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

from app.core.websocket.manager import websocket_manager
from app.utils.structured_log import get_structured_logger, LogLevel

logger = get_structured_logger(__name__, LogLevel.DEBUG)

async def test_heartbeat_task_creation():
    """测试心跳任务创建逻辑（不依赖实际连接）"""
    print("=== 心跳任务创建测试 ===")
    
    try:
        # 创建客户端
        client_name = "test_task_creation"
        client_url = "wss://localhost:8080"
        heartbeat_interval = 5  # 5秒心跳间隔，便于测试
        
        print(f"创建客户端: {client_name}")
        print(f"心跳间隔: {heartbeat_interval}秒")
        
        # 创建客户端
        client = await websocket_manager.create_client(client_name, client_url, heartbeat_interval)
        if not client:
            print("❌ 客户端创建失败")
            return False
        
        print(f"✅ 客户端创建成功: {client.name}")
        print(f"  心跳间隔配置: {client.heartbeat_interval}")
        
        # 检查心跳Channel是否存在
        heartbeat_channel_name = f"{client_name}_heartbeat"
        heartbeat_channel = websocket_manager.get_channel(heartbeat_channel_name)
        if heartbeat_channel:
            print(f"✅ 心跳Channel存在: {heartbeat_channel.name}")
        else:
            print(f"❌ 心跳Channel不存在: {heartbeat_channel_name}")
            return False
        
        # 手动启动心跳任务（模拟连接成功）
        print("手动启动心跳任务...")
        await websocket_manager._start_heartbeat_task(client_name, heartbeat_interval)
        
        # 检查心跳任务
        if client_name in websocket_manager._heartbeat_tasks:
            heartbeat_task = websocket_manager._heartbeat_tasks[client_name]
            print(f"✅ 心跳任务存在: {heartbeat_task}")
            print(f"  任务状态: {heartbeat_task.done()}")
            print(f"  任务取消状态: {heartbeat_task.cancelled()}")
        else:
            print("❌ 心跳任务不存在")
            return False
        
        # 等待心跳发送
        print(f"等待 {heartbeat_interval + 2} 秒观察心跳消息...")
        await asyncio.sleep(heartbeat_interval + 2)
        
        # 检查心跳任务状态
        if client_name in websocket_manager._heartbeat_tasks:
            heartbeat_task = websocket_manager._heartbeat_tasks[client_name]
            if not heartbeat_task.done():
                print("✅ 心跳任务仍在运行")
            else:
                print("❌ 心跳任务已结束")
                if heartbeat_task.exception():
                    print(f"  异常: {heartbeat_task.exception()}")
        else:
            print("❌ 心跳任务不存在")
        
        # 清理资源
        print("清理资源...")
        await websocket_manager.remove_client(client_name)
        print("✅ 资源清理完成")
        
        return True
        
    except Exception as e:
        print(f"❌ 测试失败: {e}")
        import traceback
        traceback.print_exc()
        return False

async def test_heartbeat_with_mock_connection():
    """测试带模拟连接的心跳功能"""
    print("\n=== 模拟连接心跳测试 ===")
    
    try:
        # 创建客户端
        client_name = "test_mock_connection"
        client_url = "wss://localhost:8080"
        heartbeat_interval = 3  # 3秒心跳间隔，便于测试
        
        print(f"创建客户端: {client_name}")
        print(f"心跳间隔: {heartbeat_interval}秒")
        
        # 创建客户端
        client = await websocket_manager.create_client(client_name, client_url, heartbeat_interval)
        if not client:
            print("❌ 客户端创建失败")
            return False
        
        print(f"✅ 客户端创建成功: {client.name}")
        
        # 模拟连接成功（直接设置连接状态）
        client._state = client._state.__class__.CONNECTED
        print("✅ 模拟连接成功")
        
        # 手动启动心跳任务
        print("启动心跳任务...")
        await websocket_manager._start_heartbeat_task(client_name, heartbeat_interval)
        
        # 检查心跳任务
        if client_name in websocket_manager._heartbeat_tasks:
            heartbeat_task = websocket_manager._heartbeat_tasks[client_name]
            print(f"✅ 心跳任务存在: {heartbeat_task}")
        else:
            print("❌ 心跳任务不存在")
            return False
        
        # 等待心跳发送
        print(f"等待 {heartbeat_interval + 2} 秒观察心跳消息...")
        await asyncio.sleep(heartbeat_interval + 2)
        
        # 检查心跳任务状态
        if client_name in websocket_manager._heartbeat_tasks:
            heartbeat_task = websocket_manager._heartbeat_tasks[client_name]
            if not heartbeat_task.done():
                print("✅ 心跳任务仍在运行")
            else:
                print("❌ 心跳任务已结束")
                if heartbeat_task.exception():
                    print(f"  异常: {heartbeat_task.exception()}")
        else:
            print("❌ 心跳任务不存在")
        
        # 清理资源
        print("清理资源...")
        await websocket_manager.remove_client(client_name)
        print("✅ 资源清理完成")
        
        return True
        
    except Exception as e:
        print(f"❌ 测试失败: {e}")
        import traceback
        traceback.print_exc()
        return False

async def main():
    """主测试函数"""
    print("开始心跳功能测试")
    print("=" * 50)
    
    # 测试1：心跳任务创建
    success1 = await test_heartbeat_task_creation()
    
    # 测试2：模拟连接心跳
    success2 = await test_heartbeat_with_mock_connection()
    
    print("\n" + "=" * 50)
    print("测试结果汇总:")
    print(f"心跳任务创建测试: {'✅ 通过' if success1 else '❌ 失败'}")
    print(f"模拟连接心跳测试: {'✅ 通过' if success2 else '❌ 失败'}")
    
    if success1 and success2:
        print("🎉 所有测试通过！")
    else:
        print("⚠️  部分测试失败，请检查日志")

if __name__ == "__main__":
    asyncio.run(main())