From 0c9eb910a601f241d03af59858593c144ca810eb Mon Sep 17 00:00:00 2001
From: Adam Goldsmith <adam@adamgoldsmith.name>
Date: Fri, 10 Mar 2023 22:16:43 -0500
Subject: [PATCH] Allow ignoring missing attributes

---
 check_home_assistant_state.py | 55 ++++++++++++++++++++++++++---------
 1 file changed, 42 insertions(+), 13 deletions(-)

diff --git a/check_home_assistant_state.py b/check_home_assistant_state.py
index aef3f00..f53f336 100755
--- a/check_home_assistant_state.py
+++ b/check_home_assistant_state.py
@@ -18,20 +18,32 @@ _log = logging.getLogger("nagiosplugin")
 
 
 class ScalarOrUnknownContext(nagiosplugin.ScalarContext):
+    """
+    Same as a :class:`nagiosplugin.ScalarContext`, but returns UNKNOWN when the
+    value is not an int or float
+    """
+
     def evaluate(self, metric: nagiosplugin.Metric, resource: nagiosplugin.Resource):
         if isinstance(metric.value, (int, float)):
             return super().evaluate(metric, resource)
         else:
-            return self.result_cls(nagiosplugin.state.Unknown, "non-scalar value", metric)
+            return self.result_cls(
+                nagiosplugin.state.Unknown, "non-scalar value", metric
+            )
 
     def performance(self, metric: nagiosplugin.Metric, resource: nagiosplugin.Resource):
         if isinstance(metric.value, (int, float)):
             return super().performance(metric, resource)
         else:
             return nagiosplugin.Performance(
-                metric.name, "U", metric.uom,
-                self.warning, self.critical,
-                metric.min, metric.max)
+                metric.name,
+                "U",
+                metric.uom,
+                self.warning,
+                self.critical,
+                metric.min,
+                metric.max,
+            )
 
 
 @dataclasses.dataclass
@@ -43,6 +55,7 @@ class Entities(nagiosplugin.Resource):
     filters: list[str]
     attribute: Optional[str]
     friendly_name: bool
+    ignore_missing: bool
     min: Optional[float] = None
     max: Optional[float] = None
 
@@ -73,7 +86,15 @@ class Entities(nagiosplugin.Resource):
                 state["attributes"].get(k) == v for k, v in self.filters
             ):
                 if self.attribute is not None:
-                    value = state["attributes"].get(self.attribute, "missing attribute")
+                    if (
+                        self.ignore_missing
+                        and self.attribute not in state["attributes"]
+                    ):
+                        continue
+                    else:
+                        value = state["attributes"].get(
+                            self.attribute, "missing attribute"
+                        )
                     uom = None
                 else:
                     if self.numeric:
@@ -165,7 +186,12 @@ class Icinga2ConfAction(argparse.Action):
             help=help,
         )
 
-    def _format_actions(self, param_prefix: str, parser: argparse.ArgumentParser, order: Optional[int] = None):
+    def _format_actions(
+        self,
+        param_prefix: str,
+        parser: argparse.ArgumentParser,
+        order: Optional[int] = None,
+    ):
         for action in parser._actions:
             if action.dest in [self.dest, "help", "version"] or isinstance(
                 action, argparse._SubParsersAction
@@ -194,7 +220,7 @@ class Icinga2ConfAction(argparse.Action):
             print(f'      description = "{action.help}"')
 
             if order is not None:
-                print(f'      order = {order}')
+                print(f"      order = {order}")
 
             print("    }")
 
@@ -211,17 +237,15 @@ class Icinga2ConfAction(argparse.Action):
             choices = parser._subparsers._group_actions[0].choices
             for subname, subparser in choices.items():
                 print(f'object CheckCommand "{command_name}_{subname}" {{')
-                print(
-                    f'  command = [ PluginContribDir + "/{command_path}" ]'
-                )
+                print(f'  command = [ PluginContribDir + "/{command_path}" ]')
 
                 print("  arguments = {")
 
                 self._format_actions(command_name, parser, order=-2)
                 print(f'    "{subname}" = {{')
-                print('      set_if = true')
-                print('      order = -1')
-                print('    }')
+                print("      set_if = true")
+                print("      order = -1")
+                print("    }")
                 self._format_actions(f"{command_name}_{subname}", subparser)
 
                 print("  }\n}")
@@ -284,6 +308,9 @@ def main():
         action="store_true",
         help="use friendly name, when available",
     )
+    common_args.add_argument(
+        "--ignore-missing", action="store_true", help="Ignore missing attributes"
+    )
 
     subparsers = argp.add_subparsers(
         title="Query type", required=True, dest="subparser_name"
@@ -352,6 +379,7 @@ def main():
                 filters=args.filter,
                 attribute=args.attribute,
                 friendly_name=args.friendly,
+                ignore_missing=args.ignore_missing,
             ),
             ScalarOrUnknownContext(args.device_class, args.warning, args.critical),
         )
@@ -365,6 +393,7 @@ def main():
                 filters=args.filter,
                 attribute=args.attribute,
                 friendly_name=args.friendly,
+                ignore_missing=args.ignore_missing,
             ),
             RegexContext(args.device_class, args.ok, args.warning, args.critical),
         )