Improve error handling for WS API trace/get (#48502)

homeassistant/components/trace/websocket_api.py

@@ -61,7 +61,14 @@ def websocket_trace_get(hass, connection, msg):
     key = (msg["domain"], msg["item_id"])
     run_id = msg["run_id"]
 
-    trace = hass.data[DATA_TRACE][key][run_id]
+    try:
+        trace = hass.data[DATA_TRACE][key][run_id]
+    except KeyError:
+        connection.send_error(
+            msg["id"], websocket_api.ERR_NOT_FOUND, "The trace could not be found"
+        )
+        return
+
     message = websocket_api.messages.result_message(msg["id"], trace)
 
     connection.send_message(json.dumps(message, cls=TraceJSONEncoder, allow_nan=False))

tests/components/trace/test_websocket_api.py

@@ -281,6 +281,25 @@ async def test_get_trace(hass, hass_ws_client, domain, prefix):
     assert response["result"] == contexts
 
 
+@pytest.mark.parametrize("domain", ["automation", "script"])
+async def test_get_invalid_trace(hass, hass_ws_client, domain):
+    """Test getting a non-existing trace."""
+    assert await async_setup_component(hass, domain, {domain: {}})
+    client = await hass_ws_client()
+    await client.send_json(
+        {
+            "id": 1,
+            "type": "trace/get",
+            "domain": domain,
+            "item_id": "sun",
+            "run_id": "invalid",
+        }
+    )
+    response = await client.receive_json()
+    assert not response["success"]
+    assert response["error"]["code"] == "not_found"
+
+
 @pytest.mark.parametrize("domain", ["automation", "script"])
 async def test_trace_overflow(hass, hass_ws_client, domain):
     """Test the number of stored traces per script or automation is limited."""