Coverage for nexios\websockets\consumers.py: 24%

1 from .base import WebSocket

2 from nexios import status

3 import typing, json

4 from .channels import Channel, ChannelBox, PayloadTypeEnum

5 from nexios import logging

6 import uuid

7

8

9 Message = typing.MutableMapping[str, typing.Any]

10

11

12 class WebSocketConsumer:

13

14 channel: typing.Optional[Channel] = None

15 middleware: typing.List[typing.Any] = []

16

17 encoding: typing.Optional[str] = None

18

19 def __init__(

20 self,

21 logging_enabled: bool = True,

22 logger: typing.Optional[logging.Logger] = None,

23 ):

24 """

25 Args:

26 logging_enabled: Whether logging is enabled for this endpoint.

27 logger: A custom logger instance. If not provided, the default logger will be used.

28 """

29 self.logging_enabled = logging_enabled

30 self.logger = logger if logger else logging.getLogger("nexios")

31

32 @classmethod

33 def as_route(cls, path: str):

34 from nexios.routing import WebsocketRoutes

35

36 """

37 Convert the WebSocketConsumer class into a Route that can be registered with the app or router.

38 """

39

40 async def handler(websocket: WebSocket, **kwargs) -> None:

41 instance = cls()

42 await instance(websocket, **kwargs)

43

44 return WebsocketRoutes(path, handler, middlewares=cls.middleware)

45

46 async def __call__(self, ws: WebSocket) -> None:

47 self.websocket = ws

48

49 self.channel = Channel(

50 websocket=self.websocket,

51 expires=3600, # Set your desired TTL for the channel

52 payload_type=(

53 PayloadTypeEnum.JSON.value

54 if self.encoding == "json"

55 else PayloadTypeEnum.TEXT.value

56 ),

57 )

58 await self.on_connect(self.websocket)

59

60 close_code = status.WS_1000_NORMAL_CLOSURE

61

62 try:

63 while True:

64 message = await self.websocket.receive()

65 if message["type"] == "websocket.receive":

66 data = await self.decode(self.websocket, message)

67 await self.on_receive(self.websocket, data)

68 elif message["type"] == "websocket.disconnect":

69 close_code = int(

70 message.get("code") or status.WS_1000_NORMAL_CLOSURE

71 )

72 break

73 except Exception as exc:

74 close_code = status.WS_1011_INTERNAL_ERROR

75 raise exc

76 finally:

77 await self.on_disconnect(self.websocket, close_code)

78

79 async def decode(self, websocket: WebSocket, message: Message) -> typing.Any:

80 if self.encoding == "text":

81 if "text" not in message:

82 await websocket.close(code=status.WS_1003_UNSUPPORTED_DATA)

83 raise RuntimeError("Expected text websocket messages, but got bytes")

84 return message["text"]

85

86 elif self.encoding == "bytes":

87 if "bytes" not in message:

88 await websocket.close(code=status.WS_1003_UNSUPPORTED_DATA)

89 raise RuntimeError("Expected bytes websocket messages, but got text")

90 return message["bytes"]

91

92 elif self.encoding == "json":

93 if message.get("text") is not None:

94 text = message["text"]

95 else:

96 text = message["bytes"].decode("utf-8")

97

98 try:

99 return json.loads(text)

100 except json.decoder.JSONDecodeError:

101 await websocket.close(code=status.WS_1003_UNSUPPORTED_DATA)

102 raise RuntimeError("Malformed JSON data received.")

103

104 assert (

105 self.encoding is None

106 ), f"Unsupported 'encoding' attribute {self.encoding}"

107 return message["text"] if message.get("text") else message["bytes"]

108

109 async def on_connect(self, websocket: WebSocket) -> None:

110 """Override to handle an incoming websocket connection"""

111 await websocket.accept()

112 if self.logging_enabled:

113 self.logger.info("New WebSocket connection established")

114

115 async def on_receive(self, websocket: WebSocket, data: typing.Any) -> None:

116 """Override to handle an incoming websocket message"""

117 if self.logging_enabled:

118 self.logger.info(f"Received message: {data}")

119

120 async def on_disconnect(self, websocket: WebSocket, close_code: int) -> None:

121 """Override to handle a disconnecting websocket"""

122 if self.logging_enabled:

123 self.logger.info(f"WebSocket disconnected with code: {close_code}")

124

125 # New Methods for Channel and Group Management

126 async def broadcast(

127 self,

128 payload: typing.Any,

129 group_name: str = "default",

130 save_history: bool = False,

131 ) -> None:

132 """

133 Broadcast a message to all channels in a group.

134 Args:

135 payload: The message payload to broadcast.

136 group_name: The name of the group to broadcast to.

137 save_history: Whether to save the message in the group's history.

138 """

139 await ChannelBox.group_send(

140 group_name=group_name, payload=payload, save_history=save_history

141 )

142 if self.logging_enabled:

143 self.logger.info(f"Broadcasted message to group '{group_name}': {payload}")

144

145 async def send_to(self, channel_id: uuid.UUID, payload: typing.Any) -> None:

146 """

147 Send a message to a specific channel by its ID.

148 Args:

149 channel_id: The UUID of the target channel.

150 payload: The message payload to send.

151 """

152 for _, channels in ChannelBox.CHANNEL_GROUPS.items():

153 for channel in channels:

154 if channel.uuid == channel_id:

155 await channel._send(payload)

156 if self.logging_enabled:

157 self.logger.info(

158 f"Sent message to channel {channel_id}: {payload}"

159 )

160 return

161 if self.logging_enabled:

162 self.logger.warning(f"Channel with ID {channel_id} not found.")

163

164 async def group(self, group_name: str) -> typing.List[Channel]:

165 """

166 Get all channels in a specific group.

167 Args:

168 group_name: The name of the group.

169 Returns:

170 A list of channels in the group.

171 """

172 channels = list(ChannelBox.CHANNEL_GROUPS.get(group_name, {}).keys())

173 if self.logging_enabled:

174 self.logger.info(

175 f"Retrieved channels in group '{group_name}': {len(channels)} channels"

176 )

177 return channels

178

179 async def join_group(self, group_name: str) -> None:

180 """

181 Add the current channel to a group.

182 Args:

183 group_name: The name of the group to join.

184 """

185 if self.channel:

186

187 await ChannelBox.add_channel_to_group(self.channel, group_name=group_name)

188 if self.logging_enabled:

189 self.logger.info(f"Channel joined group '{group_name}'")

190

191 async def leave_group(self, group_name: str) -> None:

192 """

193 Remove the current channel from a group.

194 Args:

195 group_name: The name of the group to leave.

196 """

197 if self.channel:

198 await ChannelBox.remove_channel_from_group(

199 self.channel, group_name=group_name

200 )

201 if self.logging_enabled:

202 self.logger.info(f"Channel left group '{group_name}'")