TermControlAgent/tests/test_device_event_buffer.py

#!/usr/bin/env python3
"""
设备事件缓冲和WebSocket推送功能测试
验证在用户客户端连接WebSocket时能够推送之前的监听数据
"""
import asyncio
import json
import time
from datetime import datetime
from app.core.device.manager import device_manager
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_device_event_buffer():
    """测试设备事件缓冲和WebSocket推送功能"""
    logger.info("开始设备事件缓冲和WebSocket推送功能测试")
    
    try:
        # 1. 启动设备事件推送服务
        logger.info("步骤1: 启动设备事件推送服务")
        await device_manager.start_event_pusher()
        logger.info("设备事件推送服务启动成功")
        
        # 2. 模拟一些设备事件（在WebSocket客户端连接之前）
        logger.info("步骤2: 模拟设备事件（在WebSocket客户端连接之前）")
        test_events = [
            ("device_001", "device", {"model": "Test Device 1", "manufacturer": "Test Corp"}),
            ("device_002", "device", {"model": "Test Device 2", "manufacturer": "Test Corp"}),
            ("device_003", "offline", {"model": "Test Device 3", "manufacturer": "Test Corp"}),
            ("device_001", "offline", {"model": "Test Device 1", "manufacturer": "Test Corp"}),
        ]
        
        for device_id, status, device_info in test_events:
            logger.info(f"模拟设备事件: {device_id} - {status}")
            await device_manager.handle_auto_discovered_device_event(device_id, status, device_info)
            await asyncio.sleep(0.5)  # 短暂延迟，模拟真实场景
        
        # 3. 检查缓冲的事件
        logger.info("步骤3: 检查缓冲的事件")
        # 使用公共API获取缓冲事件，而不是直接访问内部属性
        buffered_events = await device_manager.get_buffered_events("test_check")
        buffer_size = len(buffered_events)
        logger.info(f"当前缓冲事件数量: {buffer_size}")
        
        if buffer_size > 0:
            event_types = {}
            for event in buffered_events:
                event_type = event.get('type', 'unknown')
                event_types[event_type] = event_types.get(event_type, 0) + 1
            logger.info(f"缓冲事件类型分布: {event_types}")
        
        # 4. 创建WebSocket客户端
        logger.info("步骤4: 创建WebSocket客户端")
        client_name = f"test_client_{int(time.time())}"
        test_url = "ws://localhost:8080/ws"  # 使用一个测试URL
        
        try:
            # 创建客户端（这里只是测试，不实际连接）
            client = await websocket_manager.create_client(client_name, test_url, heartbeat_interval=30)
            logger.info(f"WebSocket客户端创建成功: {client_name}")
            
            # 5. 模拟客户端连接（注册设备事件推送）
            logger.info("步骤5: 模拟客户端连接，注册设备事件推送")
            await device_manager.register_websocket_client(client_name)
            logger.info(f"客户端 {client_name} 已注册设备事件推送")
            
            # 6. 获取缓冲的事件
            logger.info("步骤6: 获取缓冲的事件")
            buffered_events = await device_manager.get_buffered_events(client_name)
            logger.info(f"获取到 {len(buffered_events)} 个缓冲事件")
            
            # 7. 验证事件内容
            logger.info("步骤7: 验证事件内容")
            for i, event in enumerate(buffered_events):
                logger.info(f"事件 {i+1}: {event.get('type')} - {event.get('device_id')} - {event.get('status')}")
                logger.info(f"  时间戳: {event.get('timestamp')}")
                logger.info(f"  设备信息: {event.get('device_info')}")
            
            # 8. 模拟推送事件到WebSocket客户端
            logger.info("步骤8: 模拟推送事件到WebSocket客户端")
            success_count = 0
            for event in buffered_events:
                try:
                    # 这里只是模拟，实际推送需要真实的WebSocket连接
                    logger.info(f"模拟推送事件: {event.get('type')} - {event.get('device_id')}")
                    success_count += 1
                except Exception as e:
                    logger.error(f"推送事件失败: {e}")
            
            logger.info(f"模拟推送完成: {success_count}/{len(buffered_events)} 个事件")
            
            # 9. 注销客户端
            logger.info("步骤9: 注销客户端")
            await device_manager.unregister_websocket_client(client_name)
            logger.info(f"客户端 {client_name} 已注销")
            
        except Exception as e:
            logger.error(f"WebSocket客户端操作失败: {e}")
        
        # 10. 测试完成
        logger.info("步骤10: 测试完成")
        logger.info("设备事件缓冲和WebSocket推送功能测试完成")
        
    except Exception as e:
        logger.error(f"测试过程中发生异常: {e}")
        import traceback
        logger.error(f"异常堆栈: {traceback.format_exc()}")
    
    finally:
        # 清理资源
        logger.info("清理测试资源")
        try:
            await device_manager.stop_event_pusher()
            logger.info("设备事件推送服务已停止")
        except Exception as e:
            logger.error(f"停止设备事件推送服务失败: {e}")

async def test_real_websocket_connection():
    """测试真实的WebSocket连接和事件推送"""
    logger.info("开始真实WebSocket连接测试")
    
    try:
        # 1. 启动设备事件推送服务
        await device_manager.start_event_pusher()
        
        # 2. 模拟一些设备事件
        logger.info("模拟设备事件...")
        await device_manager.handle_auto_discovered_device_event("real_device_001", "device", {"model": "Real Device 1"})
        await device_manager.handle_auto_discovered_device_event("real_device_002", "device", {"model": "Real Device 2"})
        
        # 3. 创建并连接WebSocket客户端（使用本地测试服务器）
        client_name = f"real_test_client_{int(time.time())}"
        test_url = "ws://localhost:8080/ws"
        
        logger.info(f"创建WebSocket客户端: {client_name}")
        client = await websocket_manager.create_client(client_name, test_url)
        
        logger.info(f"连接WebSocket客户端: {client_name}")
        success = await websocket_manager.connect_client(client_name)
        
        if success:
            logger.info(f"WebSocket客户端连接成功: {client_name}")
            
            # 等待一段时间，让事件推送完成
            logger.info("等待事件推送完成...")
            await asyncio.sleep(2)
            
            # 断开连接
            logger.info(f"断开WebSocket客户端: {client_name}")
            await websocket_manager.disconnect_client(client_name)
        else:
            logger.error(f"WebSocket客户端连接失败: {client_name}")
        
    except Exception as e:
        logger.error(f"真实WebSocket连接测试失败: {e}")
    
    finally:
        await device_manager.stop_event_pusher()

if __name__ == "__main__":
    logger.info("=" * 60)
    logger.info("设备事件缓冲和WebSocket推送功能测试")
    logger.info("=" * 60)
    
    # 运行测试
    asyncio.run(test_device_event_buffer())
    
    # 可选：运行真实WebSocket连接测试
    # asyncio.run(test_real_websocket_connection())
    
    logger.info("=" * 60)
    logger.info("测试完成")
    logger.info("=" * 60)