From 3ddb67f7532b8449aca00f5149cf0bebb0ef0c2b Mon Sep 17 00:00:00 2001
From: Adam Goldsmith <adam@adamgoldsmith.name>
Date: Sat, 26 Nov 2022 19:42:49 -0500
Subject: [PATCH] Add cli arg for specifying device class to filter by

---
 check_home_assistant_state.py | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/check_home_assistant_state.py b/check_home_assistant_state.py
index 193dfd5..fc18cc6 100755
--- a/check_home_assistant_state.py
+++ b/check_home_assistant_state.py
@@ -16,6 +16,7 @@ _log = logging.getLogger("nagiosplugin")
 class Entities(nagiosplugin.Resource):
     url: str
     token: str
+    device_class: str
     min: float
     max: float
 
@@ -42,14 +43,14 @@ 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:
                 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,
                     uom=state["attributes"].get("unit_of_measurement"),
-                    context="battery",
+                    context=self.device_class,
                     min=self.min,
                     max=self.max,
                 )
@@ -66,6 +67,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",
@@ -97,6 +105,7 @@ def main():
         Entities(
             args.url,
             args.token,
+            args.device_class,
             args.min,
             args.max,
         ),