VWED_server/tests/test_websocket_module.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
WebSocket模块测试
测试WebSocket内置函数的基本功能，包括模拟WebSocket服务和客户端
"""

import asyncio
import json
import time
import threading
import sys
import os
from unittest.mock import Mock, AsyncMock, patch
from typing import Dict, Any, Set

# 添加项目根目录到Python路径
current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(current_dir)
if parent_dir not in sys.path:
    sys.path.insert(0, parent_dir)

# 导入被测试的模块
try:
    from services.online_script.built_in_modules.websocket_module import VWEDWebSocketModule
    from utils.logger import get_logger
except ImportError as e:
    print(f"导入项目模块失败: {e}")
    print("请确保从项目根目录运行此脚本")
    sys.exit(1)

logger = get_logger("tests.websocket_module")


class MockWebSocket:
    """模拟WebSocket连接"""
    
    def __init__(self, client_ip: str, client_id: str = None):
        self.client = Mock()
        self.client.host = client_ip
        self.client_id = client_id or f"client_{client_ip.replace('.', '_')}"
        self.messages = []
        self.closed = False
    
    async def send_text(self, message: str):
        """模拟发送文本消息"""
        if self.closed:
            raise Exception("WebSocket已关闭")
        self.messages.append(message)
        logger.info(f"模拟WebSocket发送消息: {message} 到 {self.client.host}")
    
    def close(self):
        """关闭连接"""
        self.closed = True


class MockConnectionManager:
    """模拟WebSocket连接管理器"""
    
    def __init__(self):
        self.active_connections: Dict[str, Set[MockWebSocket]] = {}
        self.storage_location_connections: Dict[str, Set[MockWebSocket]] = {}
    
    def add_connection(self, task_record_id: str, websocket: MockWebSocket):
        """添加活跃连接"""
        if task_record_id not in self.active_connections:
            self.active_connections[task_record_id] = set()
        self.active_connections[task_record_id].add(websocket)
    
    def add_storage_connection(self, scene_id: str, websocket: MockWebSocket):
        """添加库位状态连接"""
        if scene_id not in self.storage_location_connections:
            self.storage_location_connections[scene_id] = set()
        self.storage_location_connections[scene_id].add(websocket)
    
    async def send_personal_message(self, message: str, websocket: MockWebSocket):
        """发送个人消息"""
        await websocket.send_text(message)


class WebSocketModuleTest:
    """WebSocket模块测试类"""
    
    def __init__(self):
        self.mock_manager = MockConnectionManager()
        self.module = VWEDWebSocketModule("test_script_001")
        
        # 设置模拟管理器
        self.module._manager = self.mock_manager
    
    def setup_test_connections(self):
        """设置测试连接"""
        # 创建模拟WebSocket连接
        self.ws1 = MockWebSocket("192.168.1.100", "client1")
        self.ws2 = MockWebSocket("192.168.1.101", "client2")
        self.ws3 = MockWebSocket("192.168.1.102", "client3")
        
        # 添加到连接管理器
        self.mock_manager.add_connection("task_001", self.ws1)
        self.mock_manager.add_connection("task_002", self.ws2)
        self.mock_manager.add_storage_connection("scene_001", self.ws3)
        
        logger.info("测试连接设置完成")
    
    async def test_send_msg_by_ip(self):
        """测试根据IP发送消息"""
        logger.info("=== 测试根据IP发送消息 ===")
        
        try:
            # 测试发送消息到存在的IP
            self.module.send_msg_to_wsc_by_client_ip("测试消息1", "192.168.1.100")
            await asyncio.sleep(0.1)  # 等待异步操作完成
            
            # 检查消息是否发送
            assert len(self.ws1.messages) > 0, "消息未发送到正确的WebSocket连接"
            assert "测试消息1" in self.ws1.messages, "发送的消息内容不正确"
            
            logger.info(f"✓ 成功发送消息到IP 192.168.1.100，收到消息: {self.ws1.messages}")
            
            # 测试发送消息到不存在的IP
            try:
                self.module.send_msg_to_wsc_by_client_ip("测试消息2", "192.168.1.999")
                assert False, "应该抛出异常但没有抛出"
            except Exception as e:
                logger.info(f"✓ 正确处理不存在的IP: {e}")
            
        except Exception as e:
            logger.error(f"✗ 测试根据IP发送消息失败: {e}")
            raise
    
    async def test_send_msg_by_name(self):
        """测试根据客户端名称发送消息"""
        logger.info("=== 测试根据客户端名称发送消息 ===")
        
        try:
            # 测试发送消息到存在的客户端名称
            self.module.send_msg_to_wsc_by_client_name("测试消息3", "task_001")
            await asyncio.sleep(0.1)  # 等待异步操作完成
            
            # 检查消息是否发送
            assert len(self.ws1.messages) > 1, "消息未发送到正确的WebSocket连接"
            assert "测试消息3" in self.ws1.messages, "发送的消息内容不正确"
            
            logger.info(f"✓ 成功发送消息到客户端 task_001，收到消息: {self.ws1.messages}")
            
            # 测试发送消息到不存在的客户端名称
            try:
                self.module.send_msg_to_wsc_by_client_name("测试消息4", "task_999")
                assert False, "应该抛出异常但没有抛出"
            except Exception as e:
                logger.info(f"✓ 正确处理不存在的客户端名称: {e}")
            
        except Exception as e:
            logger.error(f"✗ 测试根据客户端名称发送消息失败: {e}")
            raise
    
    def test_get_client_ips(self):
        """测试获取客户端IP列表"""
        logger.info("=== 测试获取客户端IP列表 ===")
        
        try:
            client_ips = self.module.get_websocket_client_ip()
            
            # 验证结果
            expected_ips = ["192.168.1.100", "192.168.1.101", "192.168.1.102"]
            for ip in expected_ips:
                assert ip in client_ips, f"IP {ip} 不在返回的列表中"
            
            logger.info(f"✓ 成功获取客户端IP列表: {client_ips}")
            
        except Exception as e:
            logger.error(f"✗ 测试获取客户端IP列表失败: {e}")
            raise
    
    def test_get_client_names(self):
        """测试获取客户端名称列表"""
        logger.info("=== 测试获取客户端名称列表 ===")
        
        try:
            client_names = self.module.get_websocket_client_name()
            
            # 验证结果
            expected_names = ["task_001", "task_002", "scene_001"]
            for name in expected_names:
                assert name in client_names, f"客户端名称 {name} 不在返回的列表中"
            
            logger.info(f"✓ 成功获取客户端名称列表: {client_names}")
            
        except Exception as e:
            logger.error(f"✗ 测试获取客户端名称列表失败: {e}")
            raise
    
    async def test_error_handling(self):
        """测试错误处理"""
        logger.info("=== 测试错误处理 ===")
        
        try:
            # 测试WebSocket连接异常
            self.ws1.close()
            
            try:
                self.module.send_msg_to_wsc_by_client_ip("测试消息5", "192.168.1.100")
                await asyncio.sleep(0.1)
                logger.info("✓ 处理了WebSocket连接异常")
            except Exception as e:
                logger.info(f"✓ 正确捕获WebSocket异常: {e}")
            
        except Exception as e:
            logger.error(f"✗ 测试错误处理失败: {e}")
            raise
    
    async def run_all_tests(self):
        """运行所有测试"""
        logger.info("开始WebSocket模块测试")
        
        try:
            # 设置测试环境
            self.setup_test_connections()
            
            # 运行测试
            await self.test_send_msg_by_ip()
            await self.test_send_msg_by_name()
            self.test_get_client_ips()
            self.test_get_client_names()
            await self.test_error_handling()
            
            logger.info("✓ 所有WebSocket模块测试通过")
            
        except Exception as e:
            logger.error(f"✗ WebSocket模块测试失败: {e}")
            raise


def run_mock_websocket_server():
    """运行模拟WebSocket服务器"""
    logger.info("=== 启动模拟WebSocket服务器 ===")
    
    async def mock_server():
        try:
            # 创建测试实例
            test = WebSocketModuleTest()
            
            # 运行测试
            await test.run_all_tests()
            
            logger.info("模拟WebSocket服务器测试完成")
            
        except Exception as e:
            logger.error(f"模拟WebSocket服务器测试失败: {e}")
            raise
    
    # 运行异步测试
    asyncio.run(mock_server())


def test_websocket_integration():
    """集成测试 - 模拟完整的WebSocket通信场景"""
    logger.info("=== WebSocket集成测试 ===")
    
    async def integration_test():
        # 创建测试环境
        test = WebSocketModuleTest()
        test.setup_test_connections()
        
        # 模拟实际使用场景
        logger.info("模拟实际使用场景...")
        
        # 场景1：广播消息给所有客户端
        all_ips = test.module.get_websocket_client_ip()
        for ip in all_ips:
            try:
                test.module.send_msg_to_wsc_by_client_ip(f"广播消息给 {ip}", ip)
                await asyncio.sleep(0.05)
            except Exception as e:
                logger.error(f"发送广播消息失败: {e}")
        
        # 场景2：向特定任务发送状态更新
        test.module.send_msg_to_wsc_by_client_name("任务状态: 执行中", "task_001")
        await asyncio.sleep(0.05)
        
        # 场景3：获取连接统计信息
        client_count = len(test.module.get_websocket_client_name())
        logger.info(f"当前连接客户端数量: {client_count}")
        
        logger.info("✓ WebSocket集成测试完成")
    
    asyncio.run(integration_test())


if __name__ == "__main__":
    logger.info("开始WebSocket模块完整测试")
    
    try:
        # 运行基本功能测试
        run_mock_websocket_server()
        
        # 运行集成测试
        test_websocket_integration()
        
        logger.info("✅ 所有WebSocket模块测试成功完成")
        
    except Exception as e:
        logger.error(f"❌ WebSocket模块测试失败: {e}")
        raise