From efd40f2a51234913bb4bf2a4f2785159a7617280 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.conf |  6 +++++-
 check_home_assistant_state.py   | 13 +++++++++++--
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/check_home_assistant_state.conf b/check_home_assistant_state.conf
index a675ce1..9ad09be 100644
--- a/check_home_assistant_state.conf
+++ b/check_home_assistant_state.conf
@@ -8,7 +8,11 @@ object CheckCommand "home_assistant_state" {
     "--url" = {
       value = "$home_assistant_state_url$"
       description = "URL for Home Assistant"
-    }    
+    }
+    "--device-class" = {
+      value = "$home_assistant_state_device_class$"
+      description = "Device class of entities to monitor"
+    }
     "--warning" = {
       value = "$home_assistant_state_warning$"
       description = "Return warning if battery percentage is outside RANGE"
diff --git a/check_home_assistant_state.py b/check_home_assistant_state.py
index 193dfd5..3439b9a 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="Device class of entities to monitor",
+    )
     argp.add_argument(
         "-w",
         "--warning",
@@ -97,6 +105,7 @@ def main():
         Entities(
             args.url,
             args.token,
+            args.device_class,
             args.min,
             args.max,
         ),