From a42547c0e5b58ac49952d07edea1b0d42f2c2aa0 Mon Sep 17 00:00:00 2001
From: Paulus Schoutsen <balloob@gmail.com>
Date: Wed, 23 Feb 2022 21:15:48 -0800
Subject: [PATCH] Allow get_states to recover (#67146)

---
 .../components/websocket_api/commands.py      | 34 ++++++++++++++++++-
 .../components/websocket_api/test_commands.py | 13 +++++--
 tests/components/websocket_api/test_http.py   |  3 +-
 3 files changed, 45 insertions(+), 5 deletions(-)

diff --git a/homeassistant/components/websocket_api/commands.py b/homeassistant/components/websocket_api/commands.py
index 650013dda7f..4ed0a9ada96 100644
--- a/homeassistant/components/websocket_api/commands.py
+++ b/homeassistant/components/websocket_api/commands.py
@@ -34,6 +34,10 @@ from homeassistant.helpers.json import ExtendedJSONEncoder
 from homeassistant.helpers.service import async_get_all_descriptions
 from homeassistant.loader import IntegrationNotFound, async_get_integration
 from homeassistant.setup import DATA_SETUP_TIME, async_get_loaded_integrations
+from homeassistant.util.json import (
+    find_paths_unserializable_data,
+    format_unserializable_data,
+)
 
 from . import const, decorators, messages
 from .connection import ActiveConnection
@@ -225,7 +229,35 @@ def handle_get_states(
             if entity_perm(state.entity_id, "read")
         ]
 
-    connection.send_result(msg["id"], states)
+    # JSON serialize here so we can recover if it blows up due to the
+    # state machine containing unserializable data. This command is required
+    # to succeed for the UI to show.
+    response = messages.result_message(msg["id"], states)
+    try:
+        connection.send_message(const.JSON_DUMP(response))
+        return
+    except (ValueError, TypeError):
+        connection.logger.error(
+            "Unable to serialize to JSON. Bad data found at %s",
+            format_unserializable_data(
+                find_paths_unserializable_data(response, dump=const.JSON_DUMP)
+            ),
+        )
+    del response
+
+    # If we can't serialize, we'll filter out unserializable states
+    serialized = []
+    for state in states:
+        try:
+            serialized.append(const.JSON_DUMP(state))
+        except (ValueError, TypeError):
+            # Error is already logged above
+            pass
+
+    # We now have partially serialized states. Craft some JSON.
+    response2 = const.JSON_DUMP(messages.result_message(msg["id"], ["TO_REPLACE"]))
+    response2 = response2.replace('"TO_REPLACE"', ", ".join(serialized))
+    connection.send_message(response2)
 
 
 @decorators.websocket_command({vol.Required("type"): "get_services"})
diff --git a/tests/components/websocket_api/test_commands.py b/tests/components/websocket_api/test_commands.py
index 130870f73f0..742d9bddd38 100644
--- a/tests/components/websocket_api/test_commands.py
+++ b/tests/components/websocket_api/test_commands.py
@@ -587,13 +587,20 @@ async def test_states_filters_visible(hass, hass_admin_user, websocket_client):
 
 async def test_get_states_not_allows_nan(hass, websocket_client):
     """Test get_states command not allows NaN floats."""
-    hass.states.async_set("greeting.hello", "world", {"hello": float("NaN")})
+    hass.states.async_set("greeting.hello", "world")
+    hass.states.async_set("greeting.bad", "data", {"hello": float("NaN")})
+    hass.states.async_set("greeting.bye", "universe")
 
     await websocket_client.send_json({"id": 5, "type": "get_states"})
 
     msg = await websocket_client.receive_json()
-    assert not msg["success"]
-    assert msg["error"]["code"] == const.ERR_UNKNOWN_ERROR
+    assert msg["id"] == 5
+    assert msg["type"] == const.TYPE_RESULT
+    assert msg["success"]
+    assert msg["result"] == [
+        hass.states.get("greeting.hello").as_dict(),
+        hass.states.get("greeting.bye").as_dict(),
+    ]
 
 
 async def test_subscribe_unsubscribe_events_whitelist(
diff --git a/tests/components/websocket_api/test_http.py b/tests/components/websocket_api/test_http.py
index 336c79d22b8..c3564d2b21b 100644
--- a/tests/components/websocket_api/test_http.py
+++ b/tests/components/websocket_api/test_http.py
@@ -76,7 +76,8 @@ async def test_non_json_message(hass, websocket_client, caplog):
     msg = await websocket_client.receive_json()
     assert msg["id"] == 5
     assert msg["type"] == const.TYPE_RESULT
-    assert not msg["success"]
+    assert msg["success"]
+    assert msg["result"] == []
     assert (
         f"Unable to serialize to JSON. Bad data found at $.result[0](State: test_domain.entity).attributes.bad={bad_data}(<class 'object'>"
         in caplog.text