Coverage for /Users/eugene/Development/legion-utils/legion_utils/core.py: 55%

194 statements  

« prev     ^ index     » next       coverage.py v7.3.2, created at 2023-11-01 20:15 -0400

1from enum import IntEnum 

2from pprint import pformat 

3from re import match 

4from socket import gethostname 

5from typing import Dict, Any, Optional, Union, List 

6 

7from robotnikmq import Topic, Message, RobotnikConfig 

8from typeguard import typechecked 

9 

10 

11HOSTNAME = gethostname() 

12 

13 

14@typechecked 

15class Priority(IntEnum): 

16 INFO = 0 

17 ACTIVITY = 1 

18 WARNING = 2 

19 ERROR = 3 

20 CRITICAL = 4 

21 

22 

23INFO = {"info", "information", "i"} 

24ACTIVITY = {"activity", "a"} 

25WARNING = {"warning", "warn", "w"} 

26ERROR = {"error", "err", "e"} 

27CRITICAL = {"critical", "crit", "c"} 

28 

29 

30@typechecked 

31def valid_priority(candidate: Any) -> bool: 

32 if isinstance(candidate, Priority): 

33 return True 

34 if isinstance(candidate, str): 

35 return any( 

36 candidate.lower() in P for P in [INFO, ACTIVITY, WARNING, ERROR, CRITICAL] 

37 ) 

38 if isinstance(candidate, int): 

39 return any(candidate == i for i in Priority) 

40 return False 

41 

42 

43@typechecked 

44def priority_of(candidate: Union[str, int]) -> Priority: 

45 if isinstance(candidate, str): 

46 if candidate.lower() in INFO: 

47 return Priority.INFO 

48 if candidate.lower() in ACTIVITY: 

49 return Priority.ACTIVITY 

50 if candidate.lower() in WARNING: 

51 return Priority.WARNING 

52 if candidate.lower() in ERROR: 

53 return Priority.ERROR 

54 if candidate.lower() in CRITICAL: 

55 return Priority.CRITICAL 

56 if isinstance(candidate, int): 

57 return Priority(candidate) 

58 raise ValueError(f'"{candidate}" cannot be mapped to a valid priority') 

59 

60 

61class NotificationMsg: 

62 """This is the default message type from which all other types inherit. It 

63 has the most basic and essential properties of a Legion message. Most of 

64 the properties are Optional because they're not necessary for a message 

65 type of "Notification". You probably should not be using this type 

66 directly without setting the priorities and TTLs that you want. 

67 """ 

68 @typechecked 

69 def __init__( 

70 self, 

71 contents: Dict[str, Any], 

72 alert_key: Union[str, List[str], None] = None, 

73 desc: Optional[str] = None, 

74 ttl: Optional[int] = None, 

75 priority: Optional[Priority] = None, 

76 ): 

77 default_contents: Dict[str, Any] = {'msg_src': gethostname()} 

78 self.contents = {**default_contents, **contents} 

79 self.alert_key = ( 

80 "[" + "][".join(alert_key) + "]" 

81 if isinstance(alert_key, List) 

82 else alert_key 

83 ) 

84 self.desc = desc 

85 self.ttl = ttl 

86 self.priority = priority or Priority.INFO 

87 

88 @property 

89 def msg_src(self) -> Optional[str]: 

90 """Returns the value of the msg_src property, which is typically the 

91 hostname of the machine that generated the message, however, this 

92 value can be overridden by settings it in the contents dict. 

93 

94 Returns: 

95 Optional[str]: returns the string stored in contents['msg_src'] if 

96 its present in the contents dict or None if its not. 

97 """ 

98 return self.contents.get('msg_src') 

99 

100 @typechecked 

101 def broadcast( 

102 self, exchange: str, route: str, config: Optional[RobotnikConfig] = None 

103 ) -> None: 

104 """Broadcasts the message to a given exchange and route, with an 

105 optional Robotnik configuration which can be used to broadcast the 

106 message to a wholly different set of servers than the one that would 

107 normally be used by the RobotnikMQ library. 

108 

109 Args: 

110 exchange (str): The topic exchange to broadcast to. See: 

111 https://www.rabbitmq.com/tutorials/amqp-concepts.html#exchanges 

112 route (str): The routing key for the message. See: 

113 https://www.rabbitmq.com/tutorials/tutorial-four-python.html 

114 config (Optional[RobotnikConfig], optional): An optional configuration that can be used 

115 to publish the message to a completely 

116 different set of RabbitMQ servers than what 

117 would otherwise be used by RobotnikMQ by 

118 default. Defaults to None. 

119 """ 

120 broadcast_msg(exchange, route, self, config) 

121 

122 

123@typechecked 

124class InfoMsg(NotificationMsg): 

125 def __init__(self, contents: Dict[str, Any], ttl: Optional[int] = None): 

126 super().__init__(contents=contents, ttl=ttl, priority=Priority.INFO) 

127 

128 

129@typechecked 

130class ActivityMsg(NotificationMsg): 

131 def __init__(self, contents: Dict[str, Any], ttl: Optional[int] = None): 

132 super().__init__(contents=contents, ttl=ttl, priority=Priority.ACTIVITY) 

133 

134 

135class AlertComparison: 

136 @typechecked 

137 def __init__(self, first: "AlertMsg", second: "AlertMsg"): 

138 self.key = first.key, second.key 

139 self.desc = first.desc, second.desc 

140 self.ttl = first.ttl, second.ttl 

141 self.priority = first.priority, second.priority 

142 self.contents = first.contents, second.contents 

143 

144 @property 

145 def key_equal(self) -> bool: 

146 return self.key[0] == self.key[1] 

147 

148 @property 

149 def desc_equal(self) -> bool: 

150 return self.desc[0] == self.desc[1] 

151 

152 @property 

153 def ttl_equal(self) -> bool: 

154 return self.ttl[0] == self.ttl[1] 

155 

156 @property 

157 def priority_equal(self) -> bool: 

158 return self.priority[0] == self.priority[1] 

159 

160 @property 

161 def key_match(self) -> bool: 

162 return bool(match(*self.key) or match(self.key[1], self.key[0])) 

163 

164 @property 

165 def desc_match(self) -> bool: 

166 return bool(match(*self.desc) or match(self.desc[1], self.desc[0])) 

167 

168 

169class AlertMsg(NotificationMsg): 

170 @typechecked 

171 def __init__( 

172 self, 

173 contents: Dict[str, Any], 

174 key: Union[str, List[str]], 

175 desc: str, 

176 ttl: Optional[int] = None, 

177 priority: Optional[Priority] = None, 

178 from_msg: Optional[Message] = None, 

179 ): 

180 if priority is not None and priority < Priority.WARNING: 180 ↛ 181line 180 didn't jump to line 181, because the condition on line 180 was never true

181 raise ValueError("Alerts can only have a priority of WARNING (2) or higher") 

182 if not desc: 182 ↛ 183line 182 didn't jump to line 183, because the condition on line 182 was never true

183 raise ValueError("Alerts have to have a description") 

184 if not key: 184 ↛ 185line 184 didn't jump to line 185, because the condition on line 184 was never true

185 raise ValueError("Alerts have to have a key") 

186 super().__init__( 

187 desc=desc, 

188 priority=priority or Priority.WARNING, 

189 ttl=ttl or 30, 

190 contents=contents, 

191 ) 

192 self.alert_key: str = ( 

193 "[" + "][".join(key) + "]" if isinstance(key, List) else key 

194 ) 

195 self.msg = from_msg 

196 

197 @typechecked 

198 def compare(self, other: "AlertMsg") -> AlertComparison: 

199 return AlertComparison(self, other) 

200 

201 @property 

202 def key(self) -> str: 

203 return self.alert_key 

204 

205 @property 

206 def description(self) -> Optional[str]: 

207 return self.desc 

208 

209 @staticmethod 

210 def of(msg: Message) -> "AlertMsg": # pylint: disable=C0103 

211 if "ttl" not in msg.contents: 

212 raise ValueError( 

213 f"Message id {msg.msg_id} cannot be interpreted as an alert " 

214 f"as it does not have a TTL: {pformat(msg.to_dict())}" 

215 ) 

216 if "priority" not in msg.contents: 

217 raise ValueError( 

218 f"Message id {msg.msg_id} cannot be interpreted as an alert " 

219 f"as it does not have a priority: {pformat(msg.to_dict())}" 

220 ) 

221 if not valid_priority(msg.contents["priority"]): 

222 raise ValueError( 

223 f"Message id {msg.msg_id} cannot be interpreted as an alert " 

224 f"as it does not have a valid priority: {pformat(msg.to_dict())}" 

225 ) 

226 if not isinstance(msg.contents["ttl"], int) or msg.contents["ttl"] < 0: 

227 raise ValueError( 

228 f"Message id {msg.msg_id} cannot be interpreted as an alert " 

229 f"as it does not have a valid TTL: {pformat(msg.to_dict())}" 

230 ) 

231 if "alert_key" not in msg.contents: 

232 raise ValueError( 

233 f"Message id {msg.msg_id} cannot be interpreted as an alert " 

234 f"as it does not have an alert key: {pformat(msg.to_dict())}" 

235 ) 

236 if "description" not in msg.contents: 

237 raise ValueError( 

238 f"Message id {msg.msg_id} cannot be interpreted as an alert " 

239 f"as it does not have a description: {pformat(msg.to_dict())}" 

240 ) 

241 return AlertMsg( 

242 msg.contents, 

243 key=msg.contents["alert_key"], 

244 desc=msg.contents["description"], 

245 ttl=msg.contents["ttl"], 

246 priority=priority_of(msg.contents["priority"]), 

247 from_msg=msg, 

248 ) 

249 

250 

251class WarningMsg(AlertMsg): 

252 @typechecked 

253 def __init__( 

254 self, 

255 contents: Dict[str, Any], 

256 key: Union[str, List[str]], 

257 desc: str, 

258 ttl: Optional[int] = None, 

259 ): 

260 if not desc: 

261 raise ValueError("Warnings (alerts) have to have a description") 

262 if not key: 

263 raise ValueError("Warnings (alerts) have to have a key") 

264 super().__init__( 

265 key=key, 

266 desc=desc, 

267 priority=Priority.WARNING, 

268 ttl=ttl or 30, 

269 contents=contents, 

270 ) 

271 

272 

273class ErrorMsg(AlertMsg): 

274 @typechecked 

275 def __init__( 

276 self, 

277 contents: Dict[str, Any], 

278 key: Union[str, List[str]], 

279 desc: str, 

280 ttl: Optional[int] = None, 

281 ): 

282 if not desc: 282 ↛ 283line 282 didn't jump to line 283, because the condition on line 282 was never true

283 raise ValueError("Errors (alerts) have to have a description") 

284 if not key: 284 ↛ 285line 284 didn't jump to line 285, because the condition on line 284 was never true

285 raise ValueError("Errors (alerts) have to have a key") 

286 super().__init__( 

287 key=key, 

288 desc=desc, 

289 priority=Priority.ERROR, 

290 ttl=ttl or 30, 

291 contents=contents, 

292 ) 

293 

294 

295class CriticalMsg(AlertMsg): 

296 @typechecked 

297 def __init__( 

298 self, 

299 contents: Dict[str, Any], 

300 key: Union[str, List[str]], 

301 desc: str, 

302 ttl: Optional[int] = None, 

303 ): 

304 if not desc: 

305 raise ValueError("Critical Alerts have to have a description") 

306 if not key: 

307 raise ValueError("Critical Alerts have to have a key") 

308 super().__init__( 

309 key=key, 

310 desc=desc, 

311 priority=Priority.CRITICAL, 

312 ttl=ttl or 30, 

313 contents=contents, 

314 ) 

315 

316 

317@typechecked 

318def broadcast_msg( 

319 exchange: str, 

320 route: str, 

321 notification: NotificationMsg, 

322 config: Optional[RobotnikConfig] = None, 

323) -> None: 

324 broadcast( 

325 exchange=exchange, 

326 route=route, 

327 priority=notification.priority, 

328 contents=notification.contents, 

329 ttl=notification.ttl, 

330 description=notification.desc, 

331 alert_key=notification.alert_key, 

332 config=config, 

333 ) 

334 

335 

336@typechecked 

337def broadcast_alert_msg( 

338 exchange: str, route: str, alert: AlertMsg, config: Optional[RobotnikConfig] = None 

339) -> None: 

340 broadcast( 

341 exchange=exchange, 

342 route=route, 

343 priority=alert.priority, 

344 contents=alert.contents, 

345 ttl=alert.ttl, 

346 description=alert.desc, 

347 alert_key=alert.alert_key, 

348 config=config, 

349 ) 

350 

351 

352@typechecked 

353def broadcast( 

354 exchange: str, 

355 route: str, 

356 priority: Priority, 

357 contents: Dict[str, Any], 

358 ttl: Optional[int] = None, 

359 description: Optional[str] = None, 

360 alert_key: Optional[str] = None, 

361 config: Optional[RobotnikConfig] = None, 

362): 

363 _contents: Dict[str, Any] = {"priority": priority.value} 

364 if priority.value >= 2: 

365 assert ( 

366 description is not None 

367 ), "Alerts (e.g. WARNING, ERROR, CRITICAL) must have a description" 

368 assert ( 

369 ttl is not None 

370 ), "Alerts (e.g. WARNING, ERROR, CRITICAL) must have a ttl (to clear an alert, set the ttl to 0)" 

371 assert ( 

372 alert_key is not None 

373 ), "Alerts (e.g. WARNING, ERROR, CRITICAL) must have an alert_key" 

374 if ttl is not None: 

375 _contents["ttl"] = ttl 

376 if description is not None: 

377 _contents["description"] = description 

378 if alert_key is not None: 

379 _contents["alert_key"] = alert_key 

380 _contents.update(contents) 

381 route += f".{priority.name.lower()}" 

382 Topic(exchange=exchange, config=config).broadcast( 

383 Message.of(contents=_contents), routing_key=route 

384 ) 

385 

386 

387@typechecked 

388def broadcast_info( 

389 exchange: str, 

390 route: str, 

391 contents: Dict[str, Any], 

392 config: Optional[RobotnikConfig] = None, 

393): 

394 broadcast(exchange, route, priority=Priority.INFO, contents=contents, config=config) 

395 

396 

397@typechecked 

398def broadcast_activity( 

399 exchange: str, 

400 route: str, 

401 contents: Dict[str, Any], 

402 config: Optional[RobotnikConfig] = None, 

403): 

404 broadcast( 

405 exchange, route, priority=Priority.ACTIVITY, contents=contents, config=config 

406 ) 

407 

408 

409@typechecked 

410def broadcast_alert( 

411 exchange: str, 

412 route: str, 

413 description: str, 

414 alert_key: str, 

415 contents: Dict[str, Any], 

416 ttl: int = 30, 

417 priority: Priority = Priority.WARNING, 

418 config: Optional[RobotnikConfig] = None, 

419) -> None: 

420 broadcast( 

421 exchange, 

422 route, 

423 ttl=ttl, 

424 priority=priority, 

425 contents=contents, 

426 config=config, 

427 description=description, 

428 alert_key=alert_key, 

429 ) 

430 

431 

432@typechecked 

433def broadcast_warning( 

434 exchange: str, 

435 route: str, 

436 desc: str, 

437 alert_key: str, 

438 contents: Dict[str, Any], 

439 ttl: int = 30, 

440 config: Optional[RobotnikConfig] = None, 

441): 

442 broadcast_alert( 

443 exchange, 

444 route, 

445 description=desc, 

446 alert_key=alert_key, 

447 contents=contents, 

448 ttl=ttl, 

449 priority=Priority.WARNING, 

450 config=config, 

451 ) 

452 

453 

454@typechecked 

455def broadcast_error( 

456 exchange: str, 

457 route: str, 

458 desc: str, 

459 alert_key: str, 

460 contents: Dict[str, Any], 

461 ttl: int = 30, 

462 config: Optional[RobotnikConfig] = None, 

463): 

464 broadcast_alert( 

465 exchange, 

466 route, 

467 description=desc, 

468 alert_key=alert_key, 

469 contents=contents, 

470 ttl=ttl, 

471 priority=Priority.ERROR, 

472 config=config, 

473 ) 

474 

475 

476@typechecked 

477def broadcast_critical( 

478 exchange: str, 

479 route: str, 

480 desc: str, 

481 alert_key: str, 

482 contents: Dict[str, Any], 

483 ttl: int = 30, 

484 config: Optional[RobotnikConfig] = None, 

485): 

486 broadcast_alert( 

487 exchange, 

488 route, 

489 description=desc, 

490 alert_key=alert_key, 

491 contents=contents, 

492 ttl=ttl, 

493 priority=Priority.CRITICAL, 

494 config=config, 

495 )