Fix ESPHome update entities being loaded before device_info is available (#141704)

* Fix ESPHome update entities being loaded before device_info is available

Since we load platforms when restoring config, the update
platform could be loaded before the connection to the
device was finished which meant device_info could still
be empty. Wait until device_info is available to
load the update platform.

fixes #135906

* Apply suggestions from code review

* move comment

* Update entry_data.py

Co-authored-by: TheJulianJES <TheJulianJES@users.noreply.github.com>

---------

Co-authored-by: TheJulianJES <TheJulianJES@users.noreply.github.com>

homeassistant/components/esphome/entry_data.py

@@ -282,13 +282,16 @@ class RuntimeEntryData:
     ) -> None:
         """Distribute an update of static infos to all platforms."""
         # First, load all platforms
-        needed_platforms = set()
+        needed_platforms: set[Platform] = set()
-        if async_get_dashboard(hass):
-            needed_platforms.add(Platform.UPDATE)
 
-        if self.device_info and self.device_info.voice_assistant_feature_flags_compat(
+        if self.device_info:
-            self.api_version
+            if async_get_dashboard(hass):
-        ):
+                # Only load the update platform if the device_info is set
+                # When we restore the entry, the device_info may not be set yet
+                # and we don't want to load the update platform since it needs
+                # a complete device_info.
+                needed_platforms.add(Platform.UPDATE)
+            if self.device_info.voice_assistant_feature_flags_compat(self.api_version):
                 needed_platforms.add(Platform.BINARY_SENSOR)
                 needed_platforms.add(Platform.SELECT)
 

tests/components/esphome/test_update.py

@@ -86,26 +86,28 @@ def stub_reconnect():
 )
 async def test_update_entity(
     hass: HomeAssistant,
-    stub_reconnect,
-    mock_config_entry,
-    mock_device_info,
     mock_dashboard: dict[str, Any],
-    devices_payload,
+    devices_payload: list[dict[str, Any]],
-    expected_state,
+    expected_state: str,
-    expected_attributes,
+    expected_attributes: dict[str, Any],
+    mock_client: APIClient,
+    mock_esphome_device: Callable[
+        [APIClient, list[EntityInfo], list[UserService], list[EntityState]],
+        Awaitable[MockESPHomeDevice],
+    ],
 ) -> None:
     """Test ESPHome update entity."""
     mock_dashboard["configured"] = devices_payload
     await async_get_dashboard(hass).async_refresh()
 
-    with patch(
+    await mock_esphome_device(
-        "homeassistant.components.esphome.update.DomainData.get_entry_data",
+        mock_client=mock_client,
-        return_value=Mock(available=True, device_info=mock_device_info, info={}),
+        entity_info=[],
-    ):
+        user_service=[],
-        assert await hass.config_entries.async_setup(mock_config_entry.entry_id)
+        states=[],
-        await hass.async_block_till_done()
+    )
 
-    state = hass.states.get("update.none_firmware")
+    state = hass.states.get("update.test_firmware")
     assert state is not None
     assert state.state == expected_state
     for key, expected_value in expected_attributes.items():
@@ -130,7 +132,7 @@ async def test_update_entity(
         await hass.services.async_call(
             "update",
             "install",
-            {"entity_id": "update.none_firmware"},
+            {"entity_id": "update.test_firmware"},
             blocking=True,
         )
 
@@ -155,7 +157,7 @@ async def test_update_entity(
         await hass.services.async_call(
             "update",
             "install",
-            {"entity_id": "update.none_firmware"},
+            {"entity_id": "update.test_firmware"},
             blocking=True,
         )
 
@@ -177,7 +179,7 @@ async def test_update_entity(
         await hass.services.async_call(
             "update",
             "install",
-            {"entity_id": "update.none_firmware"},
+            {"entity_id": "update.test_firmware"},
             blocking=True,
         )
 
@@ -274,28 +276,30 @@ async def test_update_device_state_for_availability(
 
 async def test_update_entity_dashboard_not_available_startup(
     hass: HomeAssistant,
-    stub_reconnect,
+    mock_client: APIClient,
-    mock_config_entry,
+    mock_esphome_device: Callable[
-    mock_device_info,
+        [APIClient, list[EntityInfo], list[UserService], list[EntityState]],
+        Awaitable[MockESPHomeDevice],
+    ],
     mock_dashboard: dict[str, Any],
 ) -> 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, info={}),
-        ),
         patch(
             "esphome_dashboard_api.ESPHomeDashboardAPI.get_devices",
             side_effect=TimeoutError,
         ),
     ):
         await async_get_dashboard(hass).async_refresh()
-        assert await hass.config_entries.async_setup(mock_config_entry.entry_id)
+        await mock_esphome_device(
-        await hass.async_block_till_done()
+            mock_client=mock_client,
+            entity_info=[],
+            user_service=[],
+            states=[],
+        )
 
     # We have a dashboard but it is not available
-    state = hass.states.get("update.none_firmware")
+    state = hass.states.get("update.test_firmware")
     assert state is None
 
     mock_dashboard["configured"] = [
@@ -308,7 +312,7 @@ async def test_update_entity_dashboard_not_available_startup(
     await async_get_dashboard(hass).async_refresh()
     await hass.async_block_till_done()
 
-    state = hass.states.get("update.none_firmware")
+    state = hass.states.get("update.test_firmware")
     assert state.state == STATE_ON
     expected_attributes = {
         "latest_version": "2023.2.0-dev",