AirTrapMine/src/fc_network_adapter/tests/test_asyncioManager.py

import os
import sys
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

import time
import threading
import socket
import random
from pymavlink import mavutil
from pymavlink.dialects.v20 import common as mavlink2

from fc_network_adapter.asyncioManager import AsyncIOManager, Connection

class MockMAVLinkSocket:
    """模擬 pymavlink socket 行為的類，用於測試"""
    
    def __init__(self, name="mock", simulate_errors=False):
        self.name = name
        self.simulate_errors = simulate_errors
        self.is_closed = False
        self.received_msgs = []
        self.seq = 0
    
    def recv_msg(self):
        """模擬接收 MAVLink 消息"""
        if self.is_closed:
            return None
            
        # 隨機模擬網絡延遲
        time.sleep(random.uniform(0.01, 0.05))
        
        # 偶爾模擬錯誤
        if self.simulate_errors and random.random() < 0.1:
            raise Exception("Simulated network error")
            
        # 70% 的機率返回消息，30% 返回 None (模擬無數據可讀)
        if random.random() > 0.3:
            msg = self._generate_mavlink_message()
            self.received_msgs.append(msg)
            return msg
        
        return None
    
    def _generate_mavlink_message(self):
        """生成模擬 MAVLink 消息"""
        # 創建 MAVLink 實例
        mav = mavlink2.MAVLink(self)
        
        # 設置來源系統和組件
        mav.srcSystem = random.randint(1, 5)
        mav.srcComponent = random.randint(1, 5)
        
        # 創建 HEARTBEAT 消息
        msg = mavlink2.MAVLink_heartbeat_message(
            type=mavlink2.MAV_TYPE_QUADROTOR,
            autopilot=mavlink2.MAV_AUTOPILOT_GENERIC,
            base_mode=0,
            custom_mode=0,
            system_status=0,
            mavlink_version=3
        )
        
        # 打包消息
        msg.pack(mav)
        self.seq += 1
        return msg
    
    def write(self, data):
        """模擬發送數據"""
        if self.is_closed:
            raise Exception("Socket is closed")
            
        # 模擬寫入操作
        time.sleep(random.uniform(0.001, 0.01))
        
        # 偶爾模擬錯誤
        if self.simulate_errors and random.random() < 0.1:
            raise Exception("Simulated write error")
            
        return len(data)
    
    def close(self):
        """關閉連接"""
        self.is_closed = True

def create_mock_mavlink_sockets(count=3):
    """創建多個模擬 MAVLink socket"""
    sockets = []
    for i in range(count):
        # 每三個 socket 中有一個模擬錯誤
        simulate_errors = (i % 3 == 0)
        socket = MockMAVLinkSocket(f"mock_{i}", simulate_errors)
        sockets.append(socket)
    return sockets

def monitor_thread(io_manager):
    """監視 AsyncIOManager 狀態的執行緒"""
    print("Starting monitoring thread...")
    
    try:
        while True:
            # 獲取總體統計資訊
            stats = io_manager.get_stats()
            print(f"\n=== AsyncIOManager Stats ===")
            print(f"Total connections: {stats.get('connection_count', 0)}")
            print(f"Total messages received: {stats.get('rx_count', 0)}")
            print(f"Total bytes received: {stats.get('rx_bytes', 0)}")
            print(f"Total errors: {stats.get('errors', 0)}")
            
            # 獲取每個連接的緩衝區狀態
            for socket_id in io_manager.connections:
                buffer = io_manager.get_buffer(socket_id)
                conn_stats = io_manager.get_stats(socket_id)
                
                if buffer and not io_manager.connections[socket_id].is_closed:
                    print(f"Connection {socket_id}: Buffer size={buffer.size()}, " + 
                          f"Messages={conn_stats.get('rx_count', 0)}, " +
                          f"Last activity={time.time() - io_manager.connections[socket_id].last_activity:.1f}s ago")
            
            # 每秒更新一次
            time.sleep(1.0)
    except KeyboardInterrupt:
        print("Monitoring thread stopped.")

def consumer_thread(io_manager, socket_id):
    """從指定連接的緩衝區消費消息"""
    print(f"Starting consumer for socket {socket_id}...")
    
    try:
        while True:
            # 獲取緩衝區
            buffer = io_manager.get_buffer(socket_id)
            if not buffer:
                print(f"Consumer {socket_id}: Buffer not found or connection closed")
                break
                
            # 獲取所有消息
            messages = buffer.get_all()
            if messages:
                print(f"Consumer {socket_id}: Received {len(messages)} messages")
                for msg in messages:
                    # 處理每條消息 (這裡只是示例，實際應用中可能有更複雜的處理)
                    print(f"  - Message from SysID={msg.get_srcSystem()}, CompID={msg.get_srcComponent()}, Type={msg.get_type()}")
            
            # 短暫休眠，避免 CPU 空轉
            time.sleep(0.2)
    except Exception as e:
        print(f"Consumer {socket_id} stopped: {e}")

def send_command_thread(io_manager, socket_ids):
    """定期向所有連接發送命令"""
    print("Starting command sender thread...")
    
    # 創建一個 MAVLink 實例用於生成消息
    mav = mavlink2.MAVLink(None)
    mav.srcSystem = 255  # GCS system ID
    mav.srcComponent = 0
    
    try:
        count = 0
        while True:
            # 等待 2 秒
            time.sleep(2.0)
            count += 1
            
            # 創建一個命令消息
            msg = mavlink2.MAVLink_command_long_message(
                target_system=1,
                target_component=1,
                command=mavlink2.MAV_CMD_REQUEST_PROTOCOL_VERSION,
                confirmation=0,
                param1=0, param2=0, param3=0,
                param4=0, param5=0, param6=0, param7=0
            )
            
            # 向所有連接發送消息
            for socket_id in socket_ids:
                if socket_id in io_manager.connections and not io_manager.connections[socket_id].is_closed:
                    try:
                        # 打包並發送消息
                        result = io_manager.send_message(socket_id, msg)
                        print(f"Sent command to socket {socket_id}: {'Success' if result else 'Failed'}")
                    except Exception as e:
                        print(f"Error sending to socket {socket_id}: {e}")
            
            # 每 10 次迭代關閉一個連接 (用於演示關閉功能)
            if count % 10 == 0 and socket_ids:
                socket_id = socket_ids.pop(0)
                print(f"Closing connection {socket_id}...")
                io_manager.remove_connection(socket_id)
    except Exception as e:
        print(f"Command sender thread stopped: {e}")

def main():
    print("AsyncIOManager Example")
    print("======================")
    
    # 創建 AsyncIOManager 實例
    io_manager = AsyncIOManager()
    
    # 啟動 I/O 管理器
    io_manager.start()
    print("AsyncIOManager started")
    
    # 創建模擬 MAVLink socket
    mock_sockets = create_mock_mavlink_sockets(5)
    socket_ids = []
    
    # 註冊連接
    for i, socket in enumerate(mock_sockets):
        socket_id = io_manager.register_connection(socket)
        socket_ids.append(socket_id)
        print(f"Registered socket {i} with ID {socket_id}")
    
    # 啟動監視執行緒
    monitor = threading.Thread(target=monitor_thread, args=(io_manager,))
    monitor.daemon = True
    monitor.start()
    
    # 為每個連接啟動消費者執行緒
    consumers = []
    for socket_id in socket_ids:
        consumer = threading.Thread(target=consumer_thread, args=(io_manager, socket_id))
        consumer.daemon = True
        consumer.start()
        consumers.append(consumer)
    
    # 啟動命令發送執行緒
    sender = threading.Thread(target=send_command_thread, args=(io_manager, socket_ids.copy()))
    sender.daemon = True
    sender.start()
    
    print("\nPress Ctrl+C to exit...")
    try:
        # 運行一段時間然後退出
        time.sleep(30)
    except KeyboardInterrupt:
        pass
    finally:
        # 清理資源
        print("\nShutting down...")
        io_manager.stop()
        print("AsyncIOManager stopped")

if __name__ == "__main__":
    main()