From cc3ae5b19ba1ffd47dca417fd5050ee8ea9d5f55 Mon Sep 17 00:00:00 2001
From: Franck Nijhof <git@frenck.dev>
Date: Tue, 7 Feb 2023 22:15:54 +0100
Subject: [PATCH] Mark ESPHome update entity unavailable when device is offline
 (#87576)

---
 .../components/esphome/entry_data.py          |  8 ++-
 homeassistant/components/esphome/update.py    | 30 +++++++++--
 tests/components/esphome/test_update.py       | 51 +++++++++++++++++++
 3 files changed, 84 insertions(+), 5 deletions(-)

diff --git a/homeassistant/components/esphome/entry_data.py b/homeassistant/components/esphome/entry_data.py
index b7443eea211..0aed6ce43a7 100644
--- a/homeassistant/components/esphome/entry_data.py
+++ b/homeassistant/components/esphome/entry_data.py
@@ -107,6 +107,11 @@ class RuntimeEntryData:
             return self.device_info.friendly_name
         return self.name
 
+    @property
+    def signal_device_updated(self) -> str:
+        """Return the signal to listen to for core device state update."""
+        return f"esphome_{self.entry_id}_on_device_update"
+
     @property
     def signal_static_info_updated(self) -> str:
         """Return the signal to listen to for updates on static info."""
@@ -207,8 +212,7 @@ class RuntimeEntryData:
     @callback
     def async_update_device_state(self, hass: HomeAssistant) -> None:
         """Distribute an update of a core device state like availability."""
-        signal = f"esphome_{self.entry_id}_on_device_update"
-        async_dispatcher_send(hass, signal)
+        async_dispatcher_send(hass, self.signal_device_updated)
 
     async def async_load_from_store(self) -> tuple[list[EntityInfo], list[UserService]]:
         """Load the retained data from store and return de-serialized data."""
diff --git a/homeassistant/components/esphome/update.py b/homeassistant/components/esphome/update.py
index 7139c9e937f..a0d7b031336 100644
--- a/homeassistant/components/esphome/update.py
+++ b/homeassistant/components/esphome/update.py
@@ -84,7 +84,10 @@ class ESPHomeUpdateEntity(CoordinatorEntity[ESPHomeDashboard], UpdateEntity):
                 (dr.CONNECTION_NETWORK_MAC, entry_data.device_info.mac_address)
             }
         )
-        if coordinator.supports_update:
+
+        # If the device has deep sleep, we can't assume we can install updates
+        # as the ESP will not be connectable (by design).
+        if coordinator.supports_update and not self._device_info.has_deep_sleep:
             self._attr_supported_features = UpdateEntityFeature.INSTALL
 
     @property
@@ -95,8 +98,16 @@ class ESPHomeUpdateEntity(CoordinatorEntity[ESPHomeDashboard], UpdateEntity):
 
     @property
     def available(self) -> bool:
-        """Return if update is available."""
-        return super().available and self._device_info.name in self.coordinator.data
+        """Return if update is available.
+
+        During deep sleep the ESP will not be connectable (by design)
+        and thus, even when unavailable, we'll show it as available.
+        """
+        return (
+            super().available
+            and (self._entry_data.available or self._device_info.has_deep_sleep)
+            and self._device_info.name in self.coordinator.data
+        )
 
     @property
     def installed_version(self) -> str | None:
@@ -133,6 +144,19 @@ class ESPHomeUpdateEntity(CoordinatorEntity[ESPHomeDashboard], UpdateEntity):
             )
         )
 
+        @callback
+        def _on_device_update() -> None:
+            """Handle update of device state, like availability."""
+            self.async_write_ha_state()
+
+        self.async_on_remove(
+            async_dispatcher_connect(
+                self.hass,
+                self._entry_data.signal_device_updated,
+                _on_device_update,
+            )
+        )
+
     async def async_install(
         self, version: str | None, backup: bool, **kwargs: Any
     ) -> None:
diff --git a/tests/components/esphome/test_update.py b/tests/components/esphome/test_update.py
index 9cfba03be9f..b10eb089ffe 100644
--- a/tests/components/esphome/test_update.py
+++ b/tests/components/esphome/test_update.py
@@ -145,3 +145,54 @@ async def test_update_static_info(
 
     state = hass.states.get("update.none_firmware")
     assert state.state == "off"
+
+
+async def test_update_device_state_for_availability(
+    hass,
+    mock_config_entry,
+    mock_device_info,
+    mock_dashboard,
+):
+    """Test ESPHome update entity changes availability with the device."""
+    mock_dashboard["configured"] = [
+        {
+            "name": "test",
+            "current_version": "1.2.3",
+        },
+    ]
+    await async_get_dashboard(hass).async_refresh()
+
+    signal_device_updated = f"esphome_{mock_config_entry.entry_id}_on_device_update"
+    runtime_data = Mock(
+        available=True,
+        device_info=mock_device_info,
+        signal_device_updated=signal_device_updated,
+    )
+
+    with patch(
+        "homeassistant.components.esphome.update.DomainData.get_entry_data",
+        return_value=runtime_data,
+    ):
+        assert await hass.config_entries.async_forward_entry_setup(
+            mock_config_entry, "update"
+        )
+
+    state = hass.states.get("update.none_firmware")
+    assert state is not None
+    assert state.state == "on"
+
+    runtime_data.available = False
+    async_dispatcher_send(hass, signal_device_updated)
+
+    state = hass.states.get("update.none_firmware")
+    assert state.state == "unavailable"
+
+    # Deep sleep devices should still be available
+    runtime_data.device_info = dataclasses.replace(
+        runtime_data.device_info, has_deep_sleep=True
+    )
+
+    async_dispatcher_send(hass, signal_device_updated)
+
+    state = hass.states.get("update.none_firmware")
+    assert state.state == "on"