From b3fd15789ad9747ac910a27786c8fef7f718ef98 Mon Sep 17 00:00:00 2001
From: Adam Goldsmith <adam@adamgoldsmith.name>
Date: Sat, 26 Nov 2022 19:44:26 -0500
Subject: [PATCH] Add cli arg to return the value of an attribute instead of
 the state

---
 check_home_assistant_state.conf |  4 ++++
 check_home_assistant_state.py   | 17 +++++++++++++++--
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/check_home_assistant_state.conf b/check_home_assistant_state.conf
index 242f34a..087b254 100644
--- a/check_home_assistant_state.conf
+++ b/check_home_assistant_state.conf
@@ -29,6 +29,10 @@ object CheckCommand "home_assistant_state" {
       value = "$home_assistant_state_max$"
       description = "Max for performance data"
     }
+    "--attribute" = {
+      value = "$home_assistant_state_attribute$"
+      description = "Check attribute instead of value"
+    }
     "--filter" = {
       value = "$home_assistant_state_filter$"
       description = "Filter by 'attribute=value' (can be an array)"
diff --git a/check_home_assistant_state.py b/check_home_assistant_state.py
index 7fc7688..a14cb94 100755
--- a/check_home_assistant_state.py
+++ b/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
@@ -20,6 +21,7 @@ class Entities(nagiosplugin.Resource):
     min: float
     max: float
     filters: list[str]
+    attribute: Optional[str]
 
     def hass_get(self, endpoint: str) -> requests.Response:
         headers = {
@@ -47,12 +49,19 @@ class Entities(nagiosplugin.Resource):
             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,
-                    uom=state["attributes"].get("unit_of_measurement"),
+                    value,
+                    uom,
                     context=self.device_class,
                     min=self.min,
                     max=self.max,
@@ -106,6 +115,9 @@ def main():
         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",
@@ -125,6 +137,7 @@ def main():
             args.min,
             args.max,
             args.filter,
+            args.attribute,
         ),
         nagiosplugin.ScalarContext(args.device_class, args.warning, args.critical),
     )