Ensure next alarm sensor associates with device

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 {}
@@ -277,6 +287,13 @@ class SAASNextAlarmSensor(RestoreEntity):
             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."""