305 lines
11 KiB
Python
305 lines
11 KiB
Python
#!/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 |