Add cli arg to return the value of an attribute instead of the state

check_home_assistant_state.py

@@ -5,6 +5,7 @@ import dataclasses
 import logging
 import sys
 from urllib.parse import urljoin
+from typing import Optional
 
 import requests
 import nagiosplugin
@@ -16,6 +17,11 @@ _log = logging.getLogger("nagiosplugin")
 class Entities(nagiosplugin.Resource):
     url: str
     token: str
+    device_class: str
+    min: float
+    max: float
+    filters: list[str]
+    attribute: Optional[str]
 
     def hass_get(self, endpoint: str) -> requests.Response:
         headers = {
@@ -40,19 +46,33 @@ class Entities(nagiosplugin.Resource):
     def probe(self):
         response = self.hass_get("/api/states")
         for state in response:
-            if state["attributes"].get("device_class") == "battery":
+            if state["attributes"].get("device_class") == self.device_class and all(
+                state["attributes"].get(k) == v for k, v in self.filters
+            ):
+                if self.attribute is not None:
+                    value = state["attributes"].get(self.attribute, -1)
+                    uom = None
+                else:
+                    value = float(state["state"]) if state["state"].isnumeric() else -1
+                    uom = state["attributes"].get("unit_of_measurement")
+
                 if state["state"] == "unavailable":
                     _log.info(f"{state['entity_id']} unavailable")
                 yield nagiosplugin.Metric(
                     state["entity_id"],
-                    float(state["state"]) if state["state"].isnumeric() else -1,
+                    value,
-                    uom=state["attributes"].get("unit_of_measurement"),
+                    uom,
-                    min=0,
+                    context=self.device_class,
-                    max=100,
+                    min=self.min,
-                    context="battery",
+                    max=self.max,
                 )
 
 
+def key_value(arg: str):
+    k, _, v = arg.partition("=")
+    return k, v
+
+
 @nagiosplugin.guarded
 def main():
     argp = argparse.ArgumentParser(description=__doc__)
@@ -64,6 +84,13 @@ def main():
     )
     argp.add_argument("-v", "--verbose", action="count", default=0)
 
+    argp.add_argument(
+        "-d",
+        "--device-class",
+        type=str,
+        required=True,
+        help="A device class",
+    )
     argp.add_argument(
         "-w",
         "--warning",
@@ -78,12 +105,41 @@ def main():
         required=True,
         help="return critical if battery percentage is outside RANGE",
     )
+    argp.add_argument(
+        "--min",
+        type=float,
+        help="Min for performance data",
+    )
+    argp.add_argument(
+        "--max",
+        type=float,
+        help="Max for performance data",
+    )
+    argp.add_argument(
+        "-a", "--attribute", type=str, help="Check attribute instead of value"
+    )
+    argp.add_argument(
+        "-f",
+        "--filter",
+        type=key_value,
+        default=[],
+        nargs="*",
+        help="Filter by 'attribute=value'",
+    )
 
     args = argp.parse_args()
 
     check = nagiosplugin.Check(
-        Entities(args.url, args.token),
+        Entities(
-        nagiosplugin.ScalarContext("battery", args.warning, args.critical),
+            args.url,
+            args.token,
+            args.device_class,
+            args.min,
+            args.max,
+            args.filter,
+            args.attribute,
+        ),
+        nagiosplugin.ScalarContext(args.device_class, args.warning, args.critical),
     )
     check.main(args.verbose)