Fix ESPHome entities unavailable if deep sleep enabled after entry setup (home-assistant#144970)

Author: oddstr13

homeassistant/components/esphome/entity.py

@@ -223,7 +223,6 @@ def __init__(
         self._states = cast(dict[int, _StateT], entry_data.state[state_type])
         assert entry_data.device_info is not None
         device_info = entry_data.device_info
-        self._device_info = device_info
         self._on_entry_data_changed()
         self._key = entity_info.key
         self._state_type = state_type
@@ -311,6 +310,11 @@ def _on_state_update(self) -> None:
     @callback
     def _on_entry_data_changed(self) -> None:
         entry_data = self._entry_data
+        # Update the device info since it can change
+        # when the device is reconnected
+        if TYPE_CHECKING:
+            assert entry_data.device_info is not None
+        self._device_info = entry_data.device_info
         self._api_version = entry_data.api_version
         self._client = entry_data.client
         if self._device_info.has_deep_sleep:

tests/components/esphome/test_entity.py

@@ -1,13 +1,15 @@
 """Test ESPHome binary sensors."""
 
 import asyncio
+from dataclasses import asdict
 from typing import Any
 from unittest.mock import AsyncMock
 
 from aioesphomeapi import (
     APIClient,
     BinarySensorInfo,
     BinarySensorState,
+    DeviceInfo,
     SensorInfo,
     SensorState,
     build_unique_id,
@@ -665,3 +667,63 @@ async def test_entity_id_preserved_on_upgrade_when_in_storage(
     )
     state = hass.states.get("binary_sensor.user_named")
     assert state is not None
+
+
+async def test_deep_sleep_added_after_setup(
+    hass: HomeAssistant,
+    mock_client: APIClient,
+    mock_esphome_device: MockESPHomeDeviceType,
+) -> None:
+    """Test deep sleep added after setup."""
+    mock_device = await mock_esphome_device(
+        mock_client=mock_client,
+        entity_info=[
+            BinarySensorInfo(
+                object_id="test",
+                key=1,
+                name="test",
+                unique_id="test",
+            ),
+        ],
+        user_service=[],
+        states=[
+            BinarySensorState(key=1, state=True, missing_state=False),
+        ],
+        device_info={"has_deep_sleep": False},
+    )
+
+    entity_id = "binary_sensor.test_test"
+    state = hass.states.get(entity_id)
+    assert state is not None
+    assert state.state == STATE_ON
+
+    await mock_device.mock_disconnect(expected_disconnect=True)
+
+    # No deep sleep, should be unavailable
+    state = hass.states.get(entity_id)
+    assert state is not None
+    assert state.state == STATE_UNAVAILABLE
+
+    await mock_device.mock_connect()
+
+    # reconnect, should be available
+    state = hass.states.get(entity_id)
+    assert state is not None
+    assert state.state == STATE_ON
+
+    await mock_device.mock_disconnect(expected_disconnect=True)
+    new_device_info = DeviceInfo(
+        **{**asdict(mock_device.device_info), "has_deep_sleep": True}
+    )
+    mock_device.client.device_info = AsyncMock(return_value=new_device_info)
+    mock_device.device_info = new_device_info
+
+    await mock_device.mock_connect()
+
+    # Now disconnect that deep sleep is set in device info
+    await mock_device.mock_disconnect(expected_disconnect=True)
+
+    # Deep sleep, should be available
+    state = hass.states.get(entity_id)
+    assert state is not None
+    assert state.state == STATE_ON