Coverage for src\core\communication.py: 0%

1""" 

2Communication Module - 通信模块 

3 

4提供设备间的消息传递功能 

5""" 

6 

7import asyncio 

8import json 

9import logging 

10from typing import Dict, Optional, Callable, Any 

11from dataclasses import dataclass 

12from enum import Enum 

13 

14logger = logging.getLogger(__name__) 

15 

16 

17class MessageType(str, Enum): 

18 """消息类型""" 

19 TEXT = "text" 

20 SYNC = "sync" 

21 COMMAND = "command" 

22 RESPONSE = "response" 

23 HEARTBEAT = "heartbeat" 

24 

25 

26@dataclass 

27class Message: 

28 """消息""" 

29 msg_id: str 

30 msg_type: MessageType 

31 sender: str 

32 receiver: str 

33 content: Any 

34 timestamp: float 

35 

36 

37class Communication: 

38 """ 

39 通信管理器 

40 

41 负责设备间的消息发送和接收 

42 """ 

43 

44 def __init__(self, device_id: str, config=None): 

45 self.device_id = device_id 

46 self.config = config 

47 self.connections: Dict[str, asyncio.StreamWriter] = {} 

48 self.message_handlers: Dict[MessageType, Callable] = {} 

49 self._running = False 

50 self._heartbeat_task: Optional[asyncio.Task] = None 

51 

52 async def connect_to(self, peer_id: str, host: str = "localhost", port: int = 8765) -> bool: 

53 """ 

54 连接到对等节点 

55 

56 Args: 

57 peer_id: 对等节点ID 

58 host: 主机地址 

59 port: 端口 

60 

61 Returns: 

62 bool: 连接是否成功 

63 """ 

64 try: 

65 reader, writer = await asyncio.open_connection(host, port) 

66 self.connections[peer_id] = writer 

67 logger.info(f"已连接到 {peer_id}") 

68 return True 

69 except Exception as e: 

70 logger.error(f"连接失败: {e}") 

71 return False 

72 

73 async def disconnect(self): 

74 """断开所有连接""" 

75 self._running = False 

76 

77 for peer_id, writer in self.connections.items(): 

78 writer.close() 

79 await writer.wait_closed() 

80 

81 self.connections.clear() 

82 logger.info("已断开所有连接") 

83 

84 async def disconnect_peer(self, peer_id: str): 

85 """断开指定对等节点的连接""" 

86 if peer_id in self.connections: 

87 writer = self.connections[peer_id] 

88 writer.close() 

89 await writer.wait_closed() 

90 del self.connections[peer_id] 

91 logger.info(f"已断开与 {peer_id} 的连接") 

92 

93 async def send_message(self, peer_id: str, content: Any, 

94 msg_type: MessageType = MessageType.TEXT) -> bool: 

95 """ 

96 发送消息 

97 

98 Args: 

99 peer_id: 对等节点ID 

100 content: 消息内容 

101 msg_type: 消息类型 

102 

103 Returns: 

104 bool: 发送是否成功 

105 """ 

106 if peer_id not in self.connections: 

107 logger.warning(f"未连接到 {peer_id}") 

108 return False 

109 

110 try: 

111 message = Message( 

112 msg_id=f"{self.device_id}_{asyncio.get_event_loop().time()}", 

113 msg_type=msg_type, 

114 sender=self.device_id, 

115 receiver=peer_id, 

116 content=content, 

117 timestamp=asyncio.get_event_loop().time() 

118 ) 

119 

120 data = json.dumps({ 

121 "msg_id": message.msg_id, 

122 "msg_type": message.msg_type.value, 

123 "sender": message.sender, 

124 "receiver": message.receiver, 

125 "content": message.content, 

126 "timestamp": message.timestamp 

127 }) 

128 

129 writer = self.connections[peer_id] 

130 writer.write(data.encode()) 

131 await writer.drain() 

132 return True 

133 

134 except Exception as e: 

135 logger.error(f"发送消息失败: {e}") 

136 return False 

137 

138 def register_handler(self, msg_type: MessageType, handler: Callable): 

139 """注册消息处理器""" 

140 self.message_handlers[msg_type] = handler 

141 

142 async def start_heartbeat(self, interval: int = 30): 

143 """启动心跳""" 

144 self._running = True 

145 

146 async def heartbeat(): 

147 while self._running: 

148 await asyncio.sleep(interval) 

149 for peer_id in self.connections: 

150 await self.send_message(peer_id, {}, MessageType.HEARTBEAT) 

151 

152 self._heartbeat_task = asyncio.create_task(heartbeat()) 

153 

154 def get_connected_peers(self) -> list: 

155 """获取已连接的节点列表""" 

156 return list(self.connections.keys()) 

157 

158 def is_connected(self, peer_id: str) -> bool: 

159 """检查是否连接到指定节点""" 

160 return peer_id in self.connections