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

check_home_assistant_state.py

@@ -1,9 +1,11 @@
 #!/usr/bin/env python3
 
 import argparse
+import dataclasses
 import logging
 import sys
 from urllib.parse import urljoin
+from typing import Optional
 
 import requests
 import nagiosplugin
@@ -11,10 +13,15 @@ import nagiosplugin
 _log = logging.getLogger("nagiosplugin")
 
 
-class Batteries(nagiosplugin.Resource):
+@dataclasses.dataclass
-    def __init__(self, url: str, token: str):
+class Entities(nagiosplugin.Resource):
-        self.url = url
+    url: str
-        self.token = token
+    token: str
+    device_class: str
+    min: float
+    max: float
+    filters: list[str]
+    attribute: Optional[str]
 
     def hass_get(self, endpoint: str) -> requests.Response:
         headers = {
@@ -39,36 +46,36 @@ class Batteries(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__)
-    argp.add_argument(
-        "-w",
-        "--warning",
-        metavar="RANGE",
-        default="",
-        help="return warning if battery percentage is outside RANGE",
-    )
-    argp.add_argument(
-        "-c",
-        "--critical",
-        metavar="RANGE",
-        default="",
-        help="return critical if battery percentage is outside RANGE",
-    )
     argp.add_argument(
         "-t", "--token", required=True, type=str, help="API token for Home Assistant"
     )
@@ -76,11 +83,63 @@ def main():
         "-u", "--url", required=True, type=str, help="URL for Home Assistant"
     )
     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",
+        metavar="RANGE",
+        required=True,
+        help="return warning if battery percentage is outside RANGE",
+    )
+    argp.add_argument(
+        "-c",
+        "--critical",
+        metavar="RANGE",
+        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(
-        Batteries(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)