VWED_server/tests/test.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
WebSocket 脚本日志订阅测试
用于监控 VWED 脚本日志推送
"""

import asyncio
import json
from datetime import datetime
import signal
import sys

try:
    import websockets
except ImportError:
    print("❌ 缺少 websockets 依赖，请运行: pip install websockets")
    sys.exit(1)


class VWEDWebSocketLogMonitor:
    """VWED WebSocket 日志监控器"""
    
    def __init__(self, host="localhost", port=8000, client_id="monitor_test", client_type="monitor"):
        self.host = host
        self.port = port
        self.client_id = client_id
        self.client_type = client_type
        self.websocket = None
        self.running = False
        self.subscribed_scripts = set()
        
    async def connect(self):
        """连接到 WebSocket 服务器"""
        uri = f"ws://{self.host}:{self.port}/api/script/ws/script-logs?client_id={self.client_id}&client_type={self.client_type}"
        
        try:
            print(f"正在连接到 WebSocket 服务器: {uri}")
            self.websocket = await websockets.connect(uri)
            self.running = True
            print(f"✅ WebSocket 连接成功建立")
            return True
        except Exception as e:
            print(f"❌ WebSocket 连接失败: {e}")
            return False
    
    async def disconnect(self):
        """断开连接"""
        self.running = False
        if self.websocket:
            await self.websocket.close()
            print("🔌 WebSocket 连接已断开")
    
    async def subscribe_script(self, script_id):
        """订阅脚本日志"""
        if not self.websocket:
            print("WebSocket 未连接")
            return False
        
        try:
            message = {
                "type": "subscribe",
                "script_id": script_id
            }
            await self.websocket.send(json.dumps(message))
            self.subscribed_scripts.add(script_id)
            print(f"📡 已发送订阅请求: {script_id}")
            return True
        except Exception as e:
            print(f"订阅脚本失败: {e}")
            return False
    
    async def unsubscribe_script(self, script_id):
        """取消订阅脚本日志"""
        if not self.websocket:
            print("WebSocket 未连接")
            return False
        
        try:
            message = {
                "type": "unsubscribe",
                "script_id": script_id
            }
            await self.websocket.send(json.dumps(message))
            self.subscribed_scripts.discard(script_id)
            print(f" 已发送取消订阅请求: {script_id}")
            return True
        except Exception as e:
            print(f"取消订阅脚本失败: {e}")
            return False
    
    async def get_status(self):
        """获取连接状态"""
        if not self.websocket:
            print("WebSocket 未连接")
            return False
        
        try:
            message = {"type": "get_status"}
            await self.websocket.send(json.dumps(message))
            print("已发送状态查询请求")
            return True
        except Exception as e:
            print(f"获取状态失败: {e}")
            return False
    
    async def send_ping(self):
        """发送心跳"""
        if not self.websocket:
            return False
        
        try:
            message = {
                "type": "ping",
                "timestamp": datetime.now().isoformat()
            }
            await self.websocket.send(json.dumps(message))
            return True
        except Exception as e:
            print(f"❌ 发送心跳失败: {e}")
            return False
    
    def format_message(self, message):
        """格式化消息显示"""
        msg_type = message.get("type", "unknown")
        timestamp = datetime.now().strftime("%H:%M:%S")
        
        if msg_type == "welcome":
            return f"🎉 [{timestamp}] 欢迎消息: {message.get('message')} (连接ID: {message.get('connection_id')})"
        
        elif msg_type == "script_log":
            script_id = message.get("script_id", "unknown")
            level = message.get("level", "INFO")
            log_msg = message.get("message", "")
            
            # 根据日志级别选择emoji
            level_emoji = {
                "INFO": "ℹ️",
                "WARNING": "⚠️", 
                "ERROR": "❌",
                "DEBUG": "🐛"
            }
            emoji = level_emoji.get(level, "📝")
            
            return f"{emoji} [{timestamp}] [{script_id}] {level}: {log_msg}"
        
        elif msg_type == "script_status":
            script_id = message.get("script_id", "unknown")
            status = message.get("status", "unknown")
            status_msg = message.get("message", "")
            return f"🔄 [{timestamp}] [{script_id}] 状态变更: {status} - {status_msg}"
        
        elif msg_type == "function_execution":
            script_id = message.get("script_id", "unknown")
            func_name = message.get("function_name", "unknown")
            result = message.get("result", {})
            return f"⚡ [{timestamp}] [{script_id}] 函数执行: {func_name} -> {json.dumps(result, ensure_ascii=False)}"
        
        elif msg_type == "subscription_success":
            return f"✅ [{timestamp}] 订阅成功: {message.get('message')}"
        
        elif msg_type == "unsubscription_success":
            return f"✅ [{timestamp}] 取消订阅成功: {message.get('message')}"
        
        elif msg_type == "status":
            data = message.get("data", {})
            total_connections = data.get("total_connections", 0)
            total_scripts = data.get("total_script_subscriptions", 0)
            return f"📊 [{timestamp}] 连接状态: {total_connections} 个连接, {total_scripts} 个脚本订阅"
        
        elif msg_type == "error":
            return f"❌ [{timestamp}] 错误: {message.get('message')}"
        
        elif msg_type == "pong":
            return f"💓 [{timestamp}] 心跳响应"
        
        else:
            return f"❓ [{timestamp}] 未知消息类型 {msg_type}: {json.dumps(message, ensure_ascii=False)}"
    
    async def listen_messages(self):
        """监听消息"""
        if not self.websocket:
            print("❌ WebSocket 未连接")
            return
        
        try:
            while self.running:
                try:
                    # 等待消息，设置超时避免无限阻塞
                    message = await asyncio.wait_for(
                        self.websocket.recv(), 
                        timeout=1.0
                    )
                    
                    try:
                        data = json.loads(message)
                        formatted_msg = self.format_message(data)
                        print(formatted_msg)
                    except json.JSONDecodeError:
                        print(f"❌ JSON 解析失败: {message}")
                        
                except asyncio.TimeoutError:
                    # 超时是正常的，继续循环
                    continue
                except websockets.exceptions.ConnectionClosed:
                    print("🔌 WebSocket 连接已关闭")
                    break
                except Exception as e:
                    print(f"❌ 接收消息异常: {e}")
                    break
                    
        except Exception as e:
            print(f"❌ 监听消息异常: {e}")
    
    async def interactive_monitor(self):
        """交互式监控"""
        print("\n" + "="*50)
        print("VWED WebSocket 脚本日志监控器")
        print("="*50)
        print("可用命令:")
        print("  sub <script_id>  - 订阅脚本日志")
        print("  unsub <script_id> - 取消订阅脚本日志")
        print("  status           - 获取连接状态")
        print("  ping             - 发送心跳")
        print("  list             - 显示已订阅的脚本")
        print("  quit             - 退出监控")
        print("="*50)
        
        # 启动消息监听任务
        listen_task = asyncio.create_task(self.listen_messages())
        
        try:
            while self.running:
                try:
                    # 等待用户输入
                    user_input = await asyncio.get_event_loop().run_in_executor(
                        None, input, "请输入命令 (sub/unsub/status/ping/list/quit): "
                    )
                    
                    command_parts = user_input.strip().split()
                    if not command_parts:
                        continue
                    
                    command = command_parts[0].lower()
                    
                    if command == "quit":
                        print("👋 退出监控...")
                        break
                    elif command == "sub" and len(command_parts) > 1:
                        script_id = command_parts[1]
                        await self.subscribe_script(script_id)
                    elif command == "unsub" and len(command_parts) > 1:
                        script_id = command_parts[1]
                        await self.unsubscribe_script(script_id)
                    elif command == "status":
                        await self.get_status()
                    elif command == "ping":
                        await self.send_ping()
                    elif command == "list":
                        if self.subscribed_scripts:
                            print(f"📋 已订阅的脚本: {', '.join(self.subscribed_scripts)}")
                        else:
                            print("📋 当前没有订阅任何脚本")
                    else:
                        print("❌ 未知命令或参数不足")
                        
                except KeyboardInterrupt:
                    print("\n👋 收到退出信号...")
                    break
                except Exception as e:
                    print(f"❌ 处理命令异常: {e}")
        
        finally:
            self.running = False
            listen_task.cancel()
            try:
                await listen_task
            except asyncio.CancelledError:
                pass


async def main():
    """主函数"""
    monitor = VWEDWebSocketLogMonitor()
    
    # 设置信号处理
    def signal_handler(signum, _):
        print(f"\n👋 收到信号 {signum}，正在退出...")
        monitor.running = False
    
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)
    
    try:
        # 连接 WebSocket
        if await monitor.connect():
            # 开始交互式监控
            await monitor.interactive_monitor()
        else:
            print("❌ 无法连接到 WebSocket 服务器")
            return 1
    
    except Exception as e:
        print(f"❌ 监控异常: {e}")
        return 1
    
    finally:
        # 清理连接
        await monitor.disconnect()
    
    return 0


if __name__ == "__main__":
    try:
        exit_code = asyncio.run(main())
        sys.exit(exit_code)
    except KeyboardInterrupt:
        print("\n👋 程序已退出")
        sys.exit(0)