From 5d5c58338fc8b65f383daffe41e5beff146ce118 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Tue, 11 Jul 2023 11:12:24 -1000
Subject: [PATCH] Fix ESPHome deep sleep devices staying unavailable after
 unexpected disconnect (#96353)

---
 homeassistant/components/esphome/manager.py |  6 +++
 tests/components/esphome/conftest.py        | 18 ++++++-
 tests/components/esphome/test_entity.py     | 56 ++++++++++++++++++++-
 3 files changed, 78 insertions(+), 2 deletions(-)

diff --git a/homeassistant/components/esphome/manager.py b/homeassistant/components/esphome/manager.py
index b87d3ac3899..026d0315238 100644
--- a/homeassistant/components/esphome/manager.py
+++ b/homeassistant/components/esphome/manager.py
@@ -378,6 +378,12 @@ class ESPHomeManager:
             assert cli.api_version is not None
             entry_data.api_version = cli.api_version
             entry_data.available = True
+            # Reset expected disconnect flag on successful reconnect
+            # as it will be flipped to False on unexpected disconnect.
+            #
+            # We use this to determine if a deep sleep device should
+            # be marked as unavailable or not.
+            entry_data.expected_disconnect = True
             if entry_data.device_info.name:
                 assert reconnect_logic is not None, "Reconnect logic must be set"
                 reconnect_logic.name = entry_data.device_info.name
diff --git a/tests/components/esphome/conftest.py b/tests/components/esphome/conftest.py
index 1dcdc559de7..f4b3bfa3ec7 100644
--- a/tests/components/esphome/conftest.py
+++ b/tests/components/esphome/conftest.py
@@ -154,6 +154,7 @@ class MockESPHomeDevice:
         self.entry = entry
         self.state_callback: Callable[[EntityState], None]
         self.on_disconnect: Callable[[bool], None]
+        self.on_connect: Callable[[bool], None]
 
     def set_state_callback(self, state_callback: Callable[[EntityState], None]) -> None:
         """Set the state callback."""
@@ -171,6 +172,14 @@ class MockESPHomeDevice:
         """Mock disconnecting."""
         await self.on_disconnect(expected_disconnect)
 
+    def set_on_connect(self, on_connect: Callable[[], None]) -> None:
+        """Set the connect callback."""
+        self.on_connect = on_connect
+
+    async def mock_connect(self) -> None:
+        """Mock connecting."""
+        await self.on_connect()
+
 
 async def _mock_generic_device_entry(
     hass: HomeAssistant,
@@ -226,6 +235,7 @@ async def _mock_generic_device_entry(
             """Init the mock."""
             super().__init__(*args, **kwargs)
             mock_device.set_on_disconnect(kwargs["on_disconnect"])
+            mock_device.set_on_connect(kwargs["on_connect"])
             self._try_connect = self.mock_try_connect
 
         async def mock_try_connect(self):
@@ -313,9 +323,15 @@ async def mock_esphome_device(
         user_service: list[UserService],
         states: list[EntityState],
         entry: MockConfigEntry | None = None,
+        device_info: dict[str, Any] | None = None,
     ) -> MockESPHomeDevice:
         return await _mock_generic_device_entry(
-            hass, mock_client, {}, (entity_info, user_service), states, entry
+            hass,
+            mock_client,
+            device_info or {},
+            (entity_info, user_service),
+            states,
+            entry,
         )
 
     return _mock_device
diff --git a/tests/components/esphome/test_entity.py b/tests/components/esphome/test_entity.py
index 1a7d62f886b..e268d065e21 100644
--- a/tests/components/esphome/test_entity.py
+++ b/tests/components/esphome/test_entity.py
@@ -11,7 +11,7 @@ from aioesphomeapi import (
     UserService,
 )
 
-from homeassistant.const import ATTR_RESTORED, STATE_ON
+from homeassistant.const import ATTR_RESTORED, STATE_ON, STATE_UNAVAILABLE
 from homeassistant.core import HomeAssistant
 
 from .conftest import MockESPHomeDevice
@@ -130,3 +130,57 @@ async def test_entity_info_object_ids(
     )
     state = hass.states.get("binary_sensor.test_object_id_is_used")
     assert state is not None
+
+
+async def test_deep_sleep_device(
+    hass: HomeAssistant,
+    mock_client: APIClient,
+    hass_storage: dict[str, Any],
+    mock_esphome_device: Callable[
+        [APIClient, list[EntityInfo], list[UserService], list[EntityState]],
+        Awaitable[MockESPHomeDevice],
+    ],
+) -> None:
+    """Test a deep sleep device."""
+    entity_info = [
+        BinarySensorInfo(
+            object_id="mybinary_sensor",
+            key=1,
+            name="my binary_sensor",
+            unique_id="my_binary_sensor",
+        ),
+    ]
+    states = [
+        BinarySensorState(key=1, state=True, missing_state=False),
+        BinarySensorState(key=2, state=True, missing_state=False),
+    ]
+    user_service = []
+    mock_device = await mock_esphome_device(
+        mock_client=mock_client,
+        entity_info=entity_info,
+        user_service=user_service,
+        states=states,
+        device_info={"has_deep_sleep": True},
+    )
+    state = hass.states.get("binary_sensor.test_mybinary_sensor")
+    assert state is not None
+    assert state.state == STATE_ON
+
+    await mock_device.mock_disconnect(False)
+    await hass.async_block_till_done()
+    state = hass.states.get("binary_sensor.test_mybinary_sensor")
+    assert state is not None
+    assert state.state == STATE_UNAVAILABLE
+
+    await mock_device.mock_connect()
+    await hass.async_block_till_done()
+
+    state = hass.states.get("binary_sensor.test_mybinary_sensor")
+    assert state is not None
+    assert state.state == STATE_ON
+
+    await mock_device.mock_disconnect(True)
+    await hass.async_block_till_done()
+    state = hass.states.get("binary_sensor.test_mybinary_sensor")
+    assert state is not None
+    assert state.state == STATE_ON