diff --git a/homeassistant/components/recorder/core.py b/homeassistant/components/recorder/core.py index 5f598c6ce40..4866c8d536a 100644 --- a/homeassistant/components/recorder/core.py +++ b/homeassistant/components/recorder/core.py @@ -570,9 +570,11 @@ class Recorder(threading.Thread): ) @callback - def async_clear_statistics(self, statistic_ids: list[str]) -> None: + def async_clear_statistics( + self, statistic_ids: list[str], *, on_done: Callable[[], None] | None = None + ) -> None: """Clear statistics for a list of statistic_ids.""" - self.queue_task(ClearStatisticsTask(statistic_ids)) + self.queue_task(ClearStatisticsTask(on_done, statistic_ids)) @callback def async_update_statistics_metadata( diff --git a/homeassistant/components/recorder/tasks.py b/homeassistant/components/recorder/tasks.py index ce517377772..783f0a80b8e 100644 --- a/homeassistant/components/recorder/tasks.py +++ b/homeassistant/components/recorder/tasks.py @@ -60,11 +60,14 @@ class ChangeStatisticsUnitTask(RecorderTask): class ClearStatisticsTask(RecorderTask): """Object to store statistics_ids which for which to remove statistics.""" + on_done: Callable[[], None] | None statistic_ids: list[str] def run(self, instance: Recorder) -> None: """Handle the task.""" statistics.clear_statistics(instance, self.statistic_ids) + if self.on_done: + self.on_done() @dataclass(slots=True) diff --git a/homeassistant/components/recorder/websocket_api.py b/homeassistant/components/recorder/websocket_api.py index 9e4de946c0b..ac917e903df 100644 --- a/homeassistant/components/recorder/websocket_api.py +++ b/homeassistant/components/recorder/websocket_api.py @@ -49,6 +49,7 @@ from .statistics import ( ) from .util import PERIOD_SCHEMA, get_instance, resolve_period +CLEAR_STATISTICS_TIME_OUT = 10 UPDATE_STATISTICS_METADATA_TIME_OUT = 10 UNIT_SCHEMA = vol.Schema( @@ -322,8 +323,8 @@ async def ws_update_statistics_issues( vol.Required("statistic_ids"): [str], } ) -@callback -def ws_clear_statistics( +@websocket_api.async_response +async def ws_clear_statistics( hass: HomeAssistant, connection: websocket_api.ActiveConnection, msg: dict[str, Any] ) -> None: """Clear statistics for a list of statistic_ids. @@ -331,7 +332,23 @@ def ws_clear_statistics( Note: The WS call posts a job to the recorder's queue and then returns, it doesn't wait until the job is completed. """ - get_instance(hass).async_clear_statistics(msg["statistic_ids"]) + done_event = asyncio.Event() + + def clear_statistics_done() -> None: + hass.loop.call_soon_threadsafe(done_event.set) + + get_instance(hass).async_clear_statistics( + msg["statistic_ids"], on_done=clear_statistics_done + ) + try: + async with asyncio.timeout(CLEAR_STATISTICS_TIME_OUT): + await done_event.wait() + except TimeoutError: + connection.send_error( + msg["id"], websocket_api.ERR_TIMEOUT, "clear_statistics timed out" + ) + return + connection.send_result(msg["id"]) diff --git a/tests/components/recorder/test_websocket_api.py b/tests/components/recorder/test_websocket_api.py index 70ad3358430..547288d1cc3 100644 --- a/tests/components/recorder/test_websocket_api.py +++ b/tests/components/recorder/test_websocket_api.py @@ -2116,6 +2116,30 @@ async def test_clear_statistics( assert response["result"] == {"sensor.test2": expected_response["sensor.test2"]} +async def test_clear_statistics_time_out( + recorder_mock: Recorder, hass: HomeAssistant, hass_ws_client: WebSocketGenerator +) -> None: + """Test removing statistics with time-out error.""" + client = await hass_ws_client() + + with ( + patch.object(recorder.tasks.ClearStatisticsTask, "run"), + patch.object(recorder.websocket_api, "CLEAR_STATISTICS_TIME_OUT", 0), + ): + await client.send_json_auto_id( + { + "type": "recorder/clear_statistics", + "statistic_ids": ["sensor.test"], + } + ) + response = await client.receive_json() + assert not response["success"] + assert response["error"] == { + "code": "timeout", + "message": "clear_statistics timed out", + } + + @pytest.mark.parametrize( ("new_unit", "new_unit_class", "new_display_unit"), [("dogs", None, "dogs"), (None, "unitless", None), ("W", "power", "kW")],