From 03fcb81a593a647cb1a8cf607d7675344c9e8150 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Thu, 28 Dec 2023 21:36:20 -1000
Subject: [PATCH] Small speed up to compressed state diff (#106624)

---
 .../components/websocket_api/messages.py      |  2 +-
 .../components/websocket_api/test_messages.py | 44 +++++++++++++++++++
 2 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/homeassistant/components/websocket_api/messages.py b/homeassistant/components/websocket_api/messages.py
index 34ca6886b5e..3aaeff6a797 100644
--- a/homeassistant/components/websocket_api/messages.py
+++ b/homeassistant/components/websocket_api/messages.py
@@ -204,7 +204,7 @@ def _state_diff(
         for key, value in new_attributes.items():
             if old_attributes.get(key) != value:
                 additions.setdefault(COMPRESSED_STATE_ATTRIBUTES, {})[key] = value
-        if removed := set(old_attributes).difference(new_attributes):
+        if removed := old_attributes.keys() - new_attributes:
             # sets are not JSON serializable by default so we convert to list
             # here if there are any values to avoid jumping into the json_encoder_default
             # for every state diff with a removed attribute
diff --git a/tests/components/websocket_api/test_messages.py b/tests/components/websocket_api/test_messages.py
index 35ed55183d4..24387a89a29 100644
--- a/tests/components/websocket_api/test_messages.py
+++ b/tests/components/websocket_api/test_messages.py
@@ -237,6 +237,50 @@ async def test_state_diff_event(hass: HomeAssistant) -> None:
         }
     }
 
+    hass.states.async_set(
+        "light.window",
+        "green",
+        {"list_attr": ["a", "b", "c", "d"], "list_attr_2": ["a", "b"]},
+        context=new_context,
+    )
+    await hass.async_block_till_done()
+    last_state_event: Event = state_change_events[-1]
+    new_state: State = last_state_event.data["new_state"]
+    message = _state_diff_event(last_state_event)
+
+    assert message == {
+        "c": {
+            "light.window": {
+                "+": {
+                    "a": {"list_attr": ["a", "b", "c", "d"], "list_attr_2": ["a", "b"]},
+                    "lu": new_state.last_updated.timestamp(),
+                }
+            }
+        }
+    }
+
+    hass.states.async_set(
+        "light.window",
+        "green",
+        {"list_attr": ["a", "b", "c", "e"]},
+        context=new_context,
+    )
+    await hass.async_block_till_done()
+    last_state_event: Event = state_change_events[-1]
+    new_state: State = last_state_event.data["new_state"]
+    message = _state_diff_event(last_state_event)
+    assert message == {
+        "c": {
+            "light.window": {
+                "+": {
+                    "a": {"list_attr": ["a", "b", "c", "e"]},
+                    "lu": new_state.last_updated.timestamp(),
+                },
+                "-": {"a": ["list_attr_2"]},
+            }
+        }
+    }
+
 
 async def test_message_to_json(caplog: pytest.LogCaptureFixture) -> None:
     """Test we can serialize websocket messages."""