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

1 from enum import IntEnum

2 from pprint import pformat

3 from re import match

4 from socket import gethostname

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

6

7 from robotnikmq import Topic, Message, RobotnikConfig

8 from typeguard import typechecked

9

10

11 HOSTNAME = gethostname()

12

13

14@typechecked

15 class Priority(IntEnum):

16 INFO = 0

17 ACTIVITY = 1

18 WARNING = 2

19 ERROR = 3

20 CRITICAL = 4

21

22

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

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

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

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

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

28

29

30@typechecked

31 def 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

44 def 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

61 class 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

124 class 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

130 class 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

135 class 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

169 class 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

251 class 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

273 class 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

295 class 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

318 def 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

337 def 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

353 def 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

388 def 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

398 def 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

410 def 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

433 def 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

455 def 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

477 def 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 )