From f913ed343962a5266eb5173c381b46678b8240ac Mon Sep 17 00:00:00 2001
From: sudoxnym <76703581+sudoxnym@users.noreply.github.com>
Date: Sat, 28 Jun 2025 21:59:31 -0600
Subject: [PATCH] Ensure next alarm sensor associates with device

---
 custom_components/saas/sensor.py | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/custom_components/saas/sensor.py b/custom_components/saas/sensor.py
index 91abd74..9d1c36f 100644
--- a/custom_components/saas/sensor.py
+++ b/custom_components/saas/sensor.py
@@ -241,6 +241,16 @@ class SAASNextAlarmSensor(RestoreEntity):
     def state(self):
         return self._state
 
+    @property
+    def device_info(self):
+        """Return information about the device."""
+        return {
+            "identifiers": {(DOMAIN, self._name)},
+            "name": self._name,
+            "manufacturer": INTEGRATION_NAME,
+            "model": MODEL,
+        }
+
     @property
     def extra_state_attributes(self):
         return {"Label": self._label} if self._label else {}
@@ -276,6 +286,13 @@ class SAASNextAlarmSensor(RestoreEntity):
             self._hass.data[DOMAIN][self.entry_id][CONF_TOPIC],
             message_received,
         )
+
+    async def async_will_remove_from_hass(self):
+        """Run when entity will be removed from hass."""
+        self.hass.states.async_set(self.entity_id, self._state)
+        _LOGGER.info(
+            f"{datetime.now().strftime('%H:%M:%S:%f')} (Line {inspect.currentframe().f_lineno}): Saved state: {self._state} for sensor {self.name}"
+        )
 
 class SAASSoundSensor(RestoreEntity):
     """Representation of a SAAS - Sleep As Android Stats sensor for Sound Events."""