From 87c22c3ad5d447d6a90d93e956ab9e694f01c331 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Thu, 6 Apr 2023 10:32:02 -1000
Subject: [PATCH] Fix state being cleared on disconnect with deep sleep esphome
 devices (#90925)

* Fix state being cleared on disconnect with deep sleep esphome devices

fixes #90923

* fix logic
---
 homeassistant/components/esphome/__init__.py   |  9 ++++++---
 homeassistant/components/esphome/entry_data.py | 10 ++++++++--
 2 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/homeassistant/components/esphome/__init__.py b/homeassistant/components/esphome/__init__.py
index d3852c3bcc3..02a5054dd1d 100644
--- a/homeassistant/components/esphome/__init__.py
+++ b/homeassistant/components/esphome/__init__.py
@@ -345,10 +345,13 @@ async def async_setup_entry(  # noqa: C901
             disconnect_cb()
         entry_data.disconnect_callbacks = []
         entry_data.available = False
-        # Clear out the states so that we will always dispatch
+        # Mark state as stale so that we will always dispatch
         # the next state update of that type when the device reconnects
-        for state_keys in entry_data.state.values():
-            state_keys.clear()
+        entry_data.stale_state = {
+            (type(entity_state), key)
+            for state_dict in entry_data.state.values()
+            for key, entity_state in state_dict.items()
+        }
         if not hass.is_stopping:
             # Avoid marking every esphome entity as unavailable on shutdown
             # since it generates a lot of state changed events and database
diff --git a/homeassistant/components/esphome/entry_data.py b/homeassistant/components/esphome/entry_data.py
index d7f25f319ac..7a6027f946b 100644
--- a/homeassistant/components/esphome/entry_data.py
+++ b/homeassistant/components/esphome/entry_data.py
@@ -70,6 +70,10 @@ class RuntimeEntryData:
     client: APIClient
     store: Store
     state: dict[type[EntityState], dict[int, EntityState]] = field(default_factory=dict)
+    # When the disconnect callback is called, we mark all states
+    # as stale so we will always dispatch a state update when the
+    # device reconnects. This is the same format as state_subscriptions.
+    stale_state: set[tuple[type[EntityState], int]] = field(default_factory=set)
     info: dict[str, dict[int, EntityInfo]] = field(default_factory=dict)
 
     # A second list of EntityInfo objects
@@ -206,9 +210,11 @@ class RuntimeEntryData:
         """Distribute an update of state information to the target."""
         key = state.key
         state_type = type(state)
+        stale_state = self.stale_state
         current_state_by_type = self.state[state_type]
         current_state = current_state_by_type.get(key, _SENTINEL)
-        if current_state == state:
+        subscription_key = (state_type, key)
+        if current_state == state and subscription_key not in stale_state:
             _LOGGER.debug(
                 "%s: ignoring duplicate update with and key %s: %s",
                 self.name,
@@ -222,8 +228,8 @@ class RuntimeEntryData:
             key,
             state,
         )
+        stale_state.discard(subscription_key)
         current_state_by_type[key] = state
-        subscription_key = (state_type, key)
         if subscription_key in self.state_subscriptions:
             self.state_subscriptions[subscription_key]()