From d9a3884aafa562b37c1e8546e976f0272e66f789 Mon Sep 17 00:00:00 2001
From: Snuffy2 <Snuffy2@users.noreply.github.com>
Date: Mon, 28 Oct 2024 21:55:59 -0400
Subject: [PATCH] Minor sensor cleanup

---
 custom_components/opnsense/sensor.py | 45 ++++++++++++++++++----------
 1 file changed, 30 insertions(+), 15 deletions(-)

diff --git a/custom_components/opnsense/sensor.py b/custom_components/opnsense/sensor.py
index ea9f555..aef9dfd 100644
--- a/custom_components/opnsense/sensor.py
+++ b/custom_components/opnsense/sensor.py
@@ -577,7 +577,9 @@ def _handle_coordinator_update(self) -> None:
         filesystem: Mapping[str, Any] = {}
         state: Mapping[str, Any] = self.coordinator.data
         if not isinstance(state, Mapping):
-            return {}
+            self._available = False
+            self.async_write_ha_state()
+            return
         for fsystem in state.get("telemetry", {}).get("filesystems", []):
             if (
                 self.entity_description.key
@@ -611,12 +613,15 @@ def _opnsense_get_interface_property_name(self) -> str:
     def _handle_coordinator_update(self) -> None:
         state: Mapping[str, Any] = self.coordinator.data
         if not isinstance(state, Mapping):
-            return {}
+            self._available = False
+            self.async_write_ha_state()
+            return
         interface_name: str = self.entity_description.key.split(".")[1]
         interface: Mapping[str, Any] = {}
         for i_interface_name, iface in state.get("interfaces", {}).items():
             if i_interface_name == interface_name:
                 interface = iface
+                break
         if not interface:
             self._available = False
             self.async_write_ha_state()
@@ -665,11 +670,14 @@ def _handle_coordinator_update(self) -> None:
         carp_interface: Mapping[str, Any] = {}
         state: Mapping[str, Any] = self.coordinator.data
         if not isinstance(state, Mapping):
-            return {}
+            self._available = False
+            self.async_write_ha_state()
+            return
         carp_interface_name: str = self.entity_description.key.split(".")[2]
         for i_interface in state.get("carp_interfaces", []):
             if slugify(i_interface["subnet"]) == carp_interface_name:
                 carp_interface = i_interface
+                break
         if not carp_interface:
             self._available = False
             self.async_write_ha_state()
@@ -719,6 +727,7 @@ def _handle_coordinator_update(self) -> None:
         for i_gateway_name, gway in state.get("gateways", {}).items():
             if i_gateway_name == gateway_name:
                 gateway = gway
+                break
         if not gateway:
             self._available = False
             self.async_write_ha_state()
@@ -880,22 +889,26 @@ def icon(self) -> str:
 
 
 class OPNsenseTempSensor(OPNsenseSensor):
-    def _opnsense_get_temp_device(self) -> str:
-        return self.entity_description.key.split(".")[2]
-
-    def _opnsense_get_temp(self) -> Mapping[str, Any]:
+    @callback
+    def _handle_coordinator_update(self) -> None:
         state: Mapping[str, Any] = self.coordinator.data
         if not isinstance(state, Mapping):
-            return {}
-        sensor_temp_device: str = self._opnsense_get_temp_device()
-        for temp_device, temp in state.get("telemetry", {}).get("temps", {}).items():
+            self._available = False
+            self.async_write_ha_state()
+            return
+        sensor_temp_device: str = self.entity_description.key.split(".")[2]
+        temp: Mapping[str, Any] = {}
+        for temp_device, temp_temp in (
+            state.get("telemetry", {}).get("temps", {}).items()
+        ):
             if temp_device == sensor_temp_device:
-                return temp
-        return {}
+                temp = temp_temp
+                break
+        if not temp:
+            self._available = False
+            self.async_write_ha_state()
+            return
 
-    @callback
-    def _handle_coordinator_update(self) -> None:
-        temp: Mapping[str, Any] = self._opnsense_get_temp()
         try:
             self._attr_native_value = temp["temperature"]
         except (TypeError, KeyError, ZeroDivisionError):
@@ -917,6 +930,8 @@ class OPNsenseDHCPLeasesSensor(OPNsenseSensor):
     def _handle_coordinator_update(self) -> None:
         state: Mapping[str, Any] = self.coordinator.data
         if not isinstance(state, Mapping):
+            self._available = False
+            self.async_write_ha_state()
             return
         if_name: str = self.entity_description.key.split(".")[1].strip()
         # _LOGGER.debug(f"[OPNsenseDHCPLeasesSensor handle_coordinator_update] if_name: {if_name}")