Coverage for src/lexigram/admin/realtime/ws_handler_registry.py: 61%

1"""WebSocket message handler registry for real-time communication.""" 

2 

3from __future__ import annotations 

4 

5from typing import Any, Protocol 

6 

7 

8class WSMessageHandler(Protocol): 

9 """Protocol for WebSocket message handlers.""" 

10 

11 async def handle( 

12 self, 

13 websocket: Any, 

14 msg: Any, 

15 connection_id: str, 

16 manager: Any, 

17 ) -> None: 

18 """Handle the WebSocket message.""" 

19 ... 

20 

21 

22class PingHandler: 

23 """Handler for PING messages.""" 

24 

25 async def handle( 

26 self, 

27 websocket: Any, 

28 msg: Any, 

29 connection_id: str, 

30 manager: Any, 

31 ) -> None: 

32 from lexigram.admin.realtime.websocket import WSMessage, WSMessageType 

33 

34 await websocket.send_json( 

35 WSMessage( 

36 type=WSMessageType.PONG, 

37 id=msg.id, 

38 ).to_dict(), 

39 ) 

40 

41 

42class SubscribeHandler: 

43 """Handler for SUBSCRIBE messages.""" 

44 

45 async def handle( 

46 self, 

47 websocket: Any, 

48 msg: Any, 

49 connection_id: str, 

50 manager: Any, 

51 ) -> None: 

52 from lexigram.admin.realtime.websocket import WSMessage, WSMessageType 

53 

54 resources = msg.data.get("resources", []) 

55 await manager.subscribe(connection_id, resources) 

56 await websocket.send_json( 

57 WSMessage( 

58 type=WSMessageType.ACK, 

59 data={"subscribed": resources}, 

60 id=msg.id, 

61 ).to_dict(), 

62 ) 

63 

64 

65class UnsubscribeHandler: 

66 """Handler for UNSUBSCRIBE messages.""" 

67 

68 async def handle( 

69 self, 

70 websocket: Any, 

71 msg: Any, 

72 connection_id: str, 

73 manager: Any, 

74 ) -> None: 

75 from lexigram.admin.realtime.websocket import WSMessage, WSMessageType 

76 

77 resources = msg.data.get("resources", []) 

78 await manager.unsubscribe(connection_id, resources) 

79 await websocket.send_json( 

80 WSMessage( 

81 type=WSMessageType.ACK, 

82 data={"unsubscribed": resources}, 

83 id=msg.id, 

84 ).to_dict(), 

85 ) 

86 

87 

88class ActionHandler: 

89 """Handler for ACTION messages.""" 

90 

91 async def handle( 

92 self, 

93 websocket: Any, 

94 msg: Any, 

95 connection_id: str, 

96 manager: Any, 

97 ) -> None: 

98 from lexigram.admin.realtime.websocket import WebSocketHandler 

99 

100 handler = WebSocketHandler(manager) # type: ignore[call-arg] 

101 await handler._handle_action(websocket, msg) # type: ignore[attr-defined] 

102 

103 

104class WSMessageTypeRegistry: 

105 """Central registry for WebSocket message type handlers.""" 

106 

107 def __init__(self) -> None: 

108 self._handlers: dict[Any, WSMessageHandler] = {} 

109 

110 @classmethod 

111 def with_defaults(cls) -> WSMessageTypeRegistry: 

112 """Return a new instance with default message handlers registered.""" 

113 from lexigram.admin.realtime.websocket import WSMessageType 

114 

115 registry = cls() 

116 registry._handlers[WSMessageType.PING] = PingHandler() 

117 registry._handlers[WSMessageType.SUBSCRIBE] = SubscribeHandler() 

118 registry._handlers[WSMessageType.UNSUBSCRIBE] = UnsubscribeHandler() 

119 registry._handlers[WSMessageType.ACTION] = ActionHandler() 

120 return registry 

121 

122 def register(self, msg_type: Any, handler: WSMessageHandler) -> None: 

123 """Register a new message handler.""" 

124 self._handlers[msg_type] = handler 

125 

126 async def handle_message( 

127 self, 

128 msg_type: Any, 

129 websocket: Any, 

130 msg: Any, 

131 connection_id: str, 

132 manager: Any, 

133 ) -> None: 

134 """Handle a message using the appropriate handler.""" 

135 handler = self._handlers.get(msg_type) 

136 if handler: 

137 await handler.handle(websocket, msg, connection_id, manager) 

138 

139 

140_ws_message_type_registry = WSMessageTypeRegistry.with_defaults() 

141 

142 

143def get_ws_message_type_registry() -> WSMessageTypeRegistry: 

144 """Get the global WebSocket message type registry.""" 

145 return _ws_message_type_registry