From 0ba55c31e88dac1540b840815e90918e68268b9d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Odd=20Str=C3=A5b=C3=B8?= <oddstr13@openshell.no>
Date: Fri, 16 May 2025 01:57:16 +0200
Subject: [PATCH] Fix ESPHome entities unavailable if deep sleep enabled after
 entry setup (#144970)

---
 homeassistant/components/esphome/entity.py |  6 ++-
 tests/components/esphome/test_entity.py    | 62 ++++++++++++++++++++++
 2 files changed, 67 insertions(+), 1 deletion(-)

diff --git a/homeassistant/components/esphome/entity.py b/homeassistant/components/esphome/entity.py
index 7b02680afee..94c4a8ffe46 100644
--- a/homeassistant/components/esphome/entity.py
+++ b/homeassistant/components/esphome/entity.py
@@ -223,7 +223,6 @@ class EsphomeEntity(EsphomeBaseEntity, Generic[_InfoT, _StateT]):
         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 @@ class EsphomeEntity(EsphomeBaseEntity, Generic[_InfoT, _StateT]):
     @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:
diff --git a/tests/components/esphome/test_entity.py b/tests/components/esphome/test_entity.py
index ee6e6b6785f..36185efeb72 100644
--- a/tests/components/esphome/test_entity.py
+++ b/tests/components/esphome/test_entity.py
@@ -1,6 +1,7 @@
 """Test ESPHome binary sensors."""
 
 import asyncio
+from dataclasses import asdict
 from typing import Any
 from unittest.mock import AsyncMock
 
@@ -8,6 +9,7 @@ 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