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
« 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
7from robotnikmq import Topic, Message, RobotnikConfig
8from typeguard import typechecked
11HOSTNAME = gethostname()
14@typechecked
15class Priority(IntEnum):
16 INFO = 0
17 ACTIVITY = 1
18 WARNING = 2
19 ERROR = 3
20 CRITICAL = 4
23INFO = {"info", "information", "i"}
24ACTIVITY = {"activity", "a"}
25WARNING = {"warning", "warn", "w"}
26ERROR = {"error", "err", "e"}
27CRITICAL = {"critical", "crit", "c"}
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
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')
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
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.
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')
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.
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)
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)
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)
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
144 @property
145 def key_equal(self) -> bool:
146 return self.key[0] == self.key[1]
148 @property
149 def desc_equal(self) -> bool:
150 return self.desc[0] == self.desc[1]
152 @property
153 def ttl_equal(self) -> bool:
154 return self.ttl[0] == self.ttl[1]
156 @property
157 def priority_equal(self) -> bool:
158 return self.priority[0] == self.priority[1]
160 @property
161 def key_match(self) -> bool:
162 return bool(match(*self.key) or match(self.key[1], self.key[0]))
164 @property
165 def desc_match(self) -> bool:
166 return bool(match(*self.desc) or match(self.desc[1], self.desc[0]))
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
197 @typechecked
198 def compare(self, other: "AlertMsg") -> AlertComparison:
199 return AlertComparison(self, other)
201 @property
202 def key(self) -> str:
203 return self.alert_key
205 @property
206 def description(self) -> Optional[str]:
207 return self.desc
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 )
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 )
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 )
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 )
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 )
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 )
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 )
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)
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 )
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 )
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 )
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 )
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 )