From ec5f50913a19a112a59ecee90403d729fa184aa4 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Wed, 26 Apr 2023 18:41:00 +0200
Subject: [PATCH] Retry creating esphome update entities later if dashboard is
 unavailable (#92042)

---
 homeassistant/components/esphome/update.py | 43 ++++++++++++----------
 tests/components/esphome/test_update.py    | 41 +++++++++++++++++++++
 2 files changed, 65 insertions(+), 19 deletions(-)

diff --git a/homeassistant/components/esphome/update.py b/homeassistant/components/esphome/update.py
index c71ba608827..618e31024b1 100644
--- a/homeassistant/components/esphome/update.py
+++ b/homeassistant/components/esphome/update.py
@@ -13,7 +13,7 @@ from homeassistant.components.update import (
     UpdateEntityFeature,
 )
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.core import HomeAssistant, callback
+from homeassistant.core import CALLBACK_TYPE, HomeAssistant, callback
 from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity import DeviceInfo
@@ -33,35 +33,36 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up ESPHome update based on a config entry."""
-    dashboard = async_get_dashboard(hass)
-
-    if dashboard is None:
+    if (dashboard := async_get_dashboard(hass)) is None:
         return
-
     entry_data = DomainData.get(hass).get_entry_data(entry)
-    unsub = None
+    unsubs: list[CALLBACK_TYPE] = []
 
-    async def setup_update_entity() -> None:
+    @callback
+    def _async_setup_update_entity() -> None:
         """Set up the update entity."""
-        nonlocal unsub
-
+        nonlocal unsubs
+        assert dashboard is not None
         # Keep listening until device is available
-        if not entry_data.available:
+        if not entry_data.available or not dashboard.last_update_success:
             return
 
-        if unsub is not None:
-            unsub()  # type: ignore[unreachable]
+        for unsub in unsubs:
+            unsub()
+        unsubs.clear()
 
-        assert dashboard is not None
         async_add_entities([ESPHomeUpdateEntity(entry_data, dashboard)])
 
-    if entry_data.available:
-        await setup_update_entity()
+    if entry_data.available and dashboard.last_update_success:
+        _async_setup_update_entity()
         return
 
-    unsub = async_dispatcher_connect(
-        hass, entry_data.signal_device_updated, setup_update_entity
-    )
+    unsubs = [
+        async_dispatcher_connect(
+            hass, entry_data.signal_device_updated, _async_setup_update_entity
+        ),
+        dashboard.async_add_listener(_async_setup_update_entity),
+    ]
 
 
 class ESPHomeUpdateEntity(CoordinatorEntity[ESPHomeDashboard], UpdateEntity):
@@ -88,7 +89,11 @@ class ESPHomeUpdateEntity(CoordinatorEntity[ESPHomeDashboard], UpdateEntity):
 
         # 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:
+        if (
+            coordinator.last_update_success
+            and coordinator.supports_update
+            and not self._device_info.has_deep_sleep
+        ):
             self._attr_supported_features = UpdateEntityFeature.INSTALL
 
     @property
diff --git a/tests/components/esphome/test_update.py b/tests/components/esphome/test_update.py
index c8789df2777..5410af96bd7 100644
--- a/tests/components/esphome/test_update.py
+++ b/tests/components/esphome/test_update.py
@@ -1,4 +1,5 @@
 """Test ESPHome update entities."""
+import asyncio
 import dataclasses
 from unittest.mock import Mock, patch
 
@@ -197,3 +198,43 @@ async def test_update_device_state_for_availability(
 
     state = hass.states.get("update.none_firmware")
     assert state.state == "on"
+
+
+async def test_update_entity_dashboard_not_available_startup(
+    hass: HomeAssistant, mock_config_entry, mock_device_info, mock_dashboard
+) -> None:
+    """Test ESPHome update entity when dashboard is not available at startup."""
+    with patch(
+        "homeassistant.components.esphome.update.DomainData.get_entry_data",
+        return_value=Mock(available=True, device_info=mock_device_info),
+    ), patch(
+        "esphome_dashboard_api.ESPHomeDashboardAPI.get_devices",
+        side_effect=asyncio.TimeoutError,
+    ):
+        await async_get_dashboard(hass).async_refresh()
+        assert await hass.config_entries.async_forward_entry_setup(
+            mock_config_entry, "update"
+        )
+
+    state = hass.states.get("update.none_firmware")
+    assert state is None
+
+    mock_dashboard["configured"] = [
+        {
+            "name": "test",
+            "current_version": "2023.2.0-dev",
+            "configuration": "test.yaml",
+        }
+    ]
+    await async_get_dashboard(hass).async_refresh()
+    await hass.async_block_till_done()
+
+    state = hass.states.get("update.none_firmware")
+    assert state.state == "on"
+    expected_attributes = {
+        "latest_version": "2023.2.0-dev",
+        "installed_version": "1.0.0",
+        "supported_features": UpdateEntityFeature.INSTALL,
+    }
+    for key, expected_value in expected_attributes.items():
+        assert state.attributes.get(key) == expected_value