diff --git a/homeassistant/components/trace/websocket_api.py b/homeassistant/components/trace/websocket_api.py index 02d718a97ec..17f3dc7860d 100644 --- a/homeassistant/components/trace/websocket_api.py +++ b/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)) diff --git a/tests/components/trace/test_websocket_api.py b/tests/components/trace/test_websocket_api.py index c4b4f476fcc..e37e6303471 100644 --- a/tests/components/trace/test_websocket_api.py +++ b/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."""