diff --git a/homeassistant/components/statistics/config_flow.py b/homeassistant/components/statistics/config_flow.py index 773c3d1c364..145a7655b36 100644 --- a/homeassistant/components/statistics/config_flow.py +++ b/homeassistant/components/statistics/config_flow.py @@ -3,14 +3,17 @@ from __future__ import annotations from collections.abc import Mapping +from datetime import timedelta from typing import Any, cast import voluptuous as vol +from homeassistant.components import websocket_api from homeassistant.components.binary_sensor import DOMAIN as BINARY_SENSOR_DOMAIN from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN from homeassistant.const import CONF_ENTITY_ID, CONF_NAME -from homeassistant.core import split_entity_id +from homeassistant.core import HomeAssistant, callback, split_entity_id +from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers.schema_config_entry_flow import ( SchemaCommonFlowHandler, SchemaConfigFlowHandler, @@ -44,6 +47,7 @@ from .sensor import ( DEFAULT_PRECISION, STATS_BINARY_SUPPORT, STATS_NUMERIC_SUPPORT, + StatisticsSensor, ) @@ -129,12 +133,14 @@ CONFIG_FLOW = { "options": SchemaFlowFormStep( schema=DATA_SCHEMA_OPTIONS, validate_user_input=validate_options, + preview="statistics", ), } OPTIONS_FLOW = { "init": SchemaFlowFormStep( DATA_SCHEMA_OPTIONS, validate_user_input=validate_options, + preview="statistics", ), } @@ -148,3 +154,86 @@ class StatisticsConfigFlowHandler(SchemaConfigFlowHandler, domain=DOMAIN): def async_config_entry_title(self, options: Mapping[str, Any]) -> str: """Return config entry title.""" return cast(str, options[CONF_NAME]) + + @staticmethod + async def async_setup_preview(hass: HomeAssistant) -> None: + """Set up preview WS API.""" + websocket_api.async_register_command(hass, ws_start_preview) + + +@websocket_api.websocket_command( + { + vol.Required("type"): "statistics/start_preview", + vol.Required("flow_id"): str, + vol.Required("flow_type"): vol.Any("config_flow", "options_flow"), + vol.Required("user_input"): dict, + } +) +@callback +def ws_start_preview( + hass: HomeAssistant, + connection: websocket_api.ActiveConnection, + msg: dict[str, Any], +) -> None: + """Generate a preview.""" + + if msg["flow_type"] == "config_flow": + flow_status = hass.config_entries.flow.async_get(msg["flow_id"]) + flow_sets = hass.config_entries.flow._handler_progress_index.get( # noqa: SLF001 + flow_status["handler"] + ) + options = {} + assert flow_sets + for active_flow in flow_sets: + options = active_flow._common_handler.options # type: ignore [attr-defined] # noqa: SLF001 + config_entry = hass.config_entries.async_get_entry(flow_status["handler"]) + entity_id = options[CONF_ENTITY_ID] + name = options[CONF_NAME] + state_characteristic = options[CONF_STATE_CHARACTERISTIC] + else: + flow_status = hass.config_entries.options.async_get(msg["flow_id"]) + config_entry = hass.config_entries.async_get_entry(flow_status["handler"]) + if not config_entry: + raise HomeAssistantError("Config entry not found") + entity_id = config_entry.options[CONF_ENTITY_ID] + name = config_entry.options[CONF_NAME] + state_characteristic = config_entry.options[CONF_STATE_CHARACTERISTIC] + + @callback + def async_preview_updated(state: str, attributes: Mapping[str, Any]) -> None: + """Forward config entry state events to websocket.""" + connection.send_message( + websocket_api.event_message( + msg["id"], {"attributes": attributes, "state": state} + ) + ) + + sampling_size = msg["user_input"].get(CONF_SAMPLES_MAX_BUFFER_SIZE) + if sampling_size: + sampling_size = int(sampling_size) + + max_age = None + if max_age_input := msg["user_input"].get(CONF_MAX_AGE): + max_age = timedelta( + hours=max_age_input["hours"], + minutes=max_age_input["minutes"], + seconds=max_age_input["seconds"], + ) + preview_entity = StatisticsSensor( + hass, + entity_id, + name, + None, + state_characteristic, + sampling_size, + max_age, + msg["user_input"].get(CONF_KEEP_LAST_SAMPLE), + msg["user_input"].get(CONF_PRECISION), + msg["user_input"].get(CONF_PERCENTILE), + ) + preview_entity.hass = hass + + connection.send_result(msg["id"]) + connection.subscriptions[msg["id"]] = preview_entity.async_start_preview( + async_preview_updated + ) diff --git a/homeassistant/components/statistics/sensor.py b/homeassistant/components/statistics/sensor.py index ca1d75b57ed..b0a0dddd05d 100644 --- a/homeassistant/components/statistics/sensor.py +++ b/homeassistant/components/statistics/sensor.py @@ -3,7 +3,7 @@ from __future__ import annotations from collections import deque -from collections.abc import Callable +from collections.abc import Callable, Mapping import contextlib from datetime import datetime, timedelta import logging @@ -371,6 +371,29 @@ class StatisticsSensor(SensorEntity): ) self._update_listener: CALLBACK_TYPE | None = None + self._preview_callback: Callable[[str, Mapping[str, Any]], None] | None = None + + @callback + def async_start_preview( + self, + preview_callback: Callable[[str, Mapping[str, Any]], None], + ) -> CALLBACK_TYPE: + """Render a preview.""" + # abort early if there is no entity_id + # as without we can't track changes + # or either size or max_age is not set + if not self._source_entity_id or ( + self._samples_max_buffer_size is None and self._samples_max_age is None + ): + self._attr_available = False + calculated_state = self._async_calculate_state() + preview_callback(calculated_state.state, calculated_state.attributes) + return self._call_on_remove_callbacks + + self._preview_callback = preview_callback + + self._async_stats_sensor_startup(self.hass) + return self._call_on_remove_callbacks @callback def _async_stats_sensor_state_listener( @@ -382,7 +405,13 @@ class StatisticsSensor(SensorEntity): return self._add_state_to_queue(new_state) self._async_purge_update_and_schedule() - self.async_write_ha_state() + + if self._preview_callback: + calculated_state = self._async_calculate_state() + self._preview_callback(calculated_state.state, calculated_state.attributes) + # only write state to the state machine if we are not in preview mode + if not self._preview_callback: + self.async_write_ha_state() @callback def _async_stats_sensor_startup(self, _: HomeAssistant) -> None: @@ -604,7 +633,9 @@ class StatisticsSensor(SensorEntity): _LOGGER.debug("%s: executing scheduled update", self.entity_id) self._async_cancel_update_listener() self._async_purge_update_and_schedule() - self.async_write_ha_state() + # only write state to the state machine if we are not in preview mode + if not self._preview_callback: + self.async_write_ha_state() def _fetch_states_from_database(self) -> list[State]: """Fetch the states from the database.""" @@ -648,7 +679,13 @@ class StatisticsSensor(SensorEntity): self._add_state_to_queue(state) self._async_purge_update_and_schedule() - self.async_write_ha_state() + + # only write state to the state machine if we are not in preview mode + if self._preview_callback: + calculated_state = self._async_calculate_state() + self._preview_callback(calculated_state.state, calculated_state.attributes) + else: + self.async_write_ha_state() _LOGGER.debug("%s: initializing from database completed", self.entity_id) def _update_attributes(self) -> None: diff --git a/tests/components/statistics/snapshots/test_config_flow.ambr b/tests/components/statistics/snapshots/test_config_flow.ambr new file mode 100644 index 00000000000..8d274cd86c6 --- /dev/null +++ b/tests/components/statistics/snapshots/test_config_flow.ambr @@ -0,0 +1,49 @@ +# serializer version: 1 +# name: test_config_flow_preview_success[missing_size_and_age] + dict({ + 'attributes': dict({ + 'friendly_name': 'Statistical characteristic', + 'icon': 'mdi:calculator', + 'state_class': 'measurement', + }), + 'state': 'unavailable', + }) +# --- +# name: test_config_flow_preview_success[success] + dict({ + 'attributes': dict({ + 'buffer_usage_ratio': 0.1, + 'friendly_name': 'Statistical characteristic', + 'icon': 'mdi:calculator', + 'source_value_valid': True, + 'state_class': 'measurement', + }), + 'state': '16.0', + }) +# --- +# name: test_options_flow_preview + dict({ + 'attributes': dict({ + 'age_coverage_ratio': 0.0, + 'buffer_usage_ratio': 0.05, + 'friendly_name': 'Statistical characteristic', + 'icon': 'mdi:calculator', + 'source_value_valid': True, + 'state_class': 'measurement', + }), + 'state': '16.0', + }) +# --- +# name: test_options_flow_preview[updated] + dict({ + 'attributes': dict({ + 'age_coverage_ratio': 0.0, + 'buffer_usage_ratio': 0.1, + 'friendly_name': 'Statistical characteristic', + 'icon': 'mdi:calculator', + 'source_value_valid': True, + 'state_class': 'measurement', + }), + 'state': '20.0', + }) +# --- diff --git a/tests/components/statistics/test_config_flow.py b/tests/components/statistics/test_config_flow.py index 7c9ed5bed47..77ccba5ba4c 100644 --- a/tests/components/statistics/test_config_flow.py +++ b/tests/components/statistics/test_config_flow.py @@ -4,7 +4,11 @@ from __future__ import annotations from unittest.mock import AsyncMock +import pytest +from syrupy import SnapshotAssertion + from homeassistant import config_entries +from homeassistant.components.recorder import Recorder from homeassistant.components.statistics import DOMAIN from homeassistant.components.statistics.sensor import ( CONF_KEEP_LAST_SAMPLE, @@ -16,12 +20,14 @@ from homeassistant.components.statistics.sensor import ( DEFAULT_NAME, STAT_AVERAGE_LINEAR, STAT_COUNT, + STAT_VALUE_MAX, ) from homeassistant.const import CONF_ENTITY_ID, CONF_NAME from homeassistant.core import HomeAssistant from homeassistant.data_entry_flow import FlowResultType from tests.common import MockConfigEntry +from tests.typing import WebSocketGenerator async def test_form_sensor(hass: HomeAssistant, mock_setup_entry: AsyncMock) -> None: @@ -271,3 +277,204 @@ async def test_entry_already_exist( assert result["type"] is FlowResultType.ABORT assert result["reason"] == "already_configured" + + +@pytest.mark.parametrize( + "user_input", + [ + ( + { + CONF_SAMPLES_MAX_BUFFER_SIZE: 10.0, + CONF_KEEP_LAST_SAMPLE: False, + CONF_PERCENTILE: 50, + CONF_PRECISION: 2, + } + ), + ( + { + CONF_KEEP_LAST_SAMPLE: False, + CONF_PERCENTILE: 50, + CONF_PRECISION: 2, + } + ), + ], + ids=("success", "missing_size_and_age"), +) +async def test_config_flow_preview_success( + recorder_mock: Recorder, + hass: HomeAssistant, + hass_ws_client: WebSocketGenerator, + user_input: str, + snapshot: SnapshotAssertion, +) -> None: + """Test the config flow preview.""" + client = await hass_ws_client(hass) + + # add state for the tests + hass.states.async_set("sensor.test_monitored", "16") + + result = await hass.config_entries.flow.async_init( + DOMAIN, context={"source": config_entries.SOURCE_USER} + ) + assert result["type"] == FlowResultType.FORM + assert result["step_id"] == "user" + assert result["errors"] is None + + result = await hass.config_entries.flow.async_configure( + result["flow_id"], + { + CONF_NAME: DEFAULT_NAME, + CONF_ENTITY_ID: "sensor.test_monitored", + }, + ) + await hass.async_block_till_done() + result = await hass.config_entries.flow.async_configure( + result["flow_id"], + { + CONF_STATE_CHARACTERISTIC: STAT_VALUE_MAX, + }, + ) + await hass.async_block_till_done() + assert result["type"] == FlowResultType.FORM + assert result["step_id"] == "options" + assert result["errors"] is None + assert result["preview"] == "statistics" + + await client.send_json_auto_id( + { + "type": "statistics/start_preview", + "flow_id": result["flow_id"], + "flow_type": "config_flow", + "user_input": user_input, + } + ) + msg = await client.receive_json() + assert msg["success"] + assert msg["result"] is None + + msg = await client.receive_json() + assert msg["event"] == snapshot + assert len(hass.states.async_all()) == 1 + + +async def test_options_flow_preview( + recorder_mock: Recorder, + hass: HomeAssistant, + hass_ws_client: WebSocketGenerator, + snapshot: SnapshotAssertion, +) -> None: + """Test the options flow preview.""" + client = await hass_ws_client(hass) + + # add state for the tests + hass.states.async_set("sensor.test_monitored", "16") + + # Setup the config entry + config_entry = MockConfigEntry( + data={}, + domain=DOMAIN, + options={ + CONF_NAME: DEFAULT_NAME, + CONF_ENTITY_ID: "sensor.test_monitored", + CONF_STATE_CHARACTERISTIC: STAT_VALUE_MAX, + CONF_SAMPLES_MAX_BUFFER_SIZE: 20.0, + CONF_MAX_AGE: {"hours": 8, "minutes": 0, "seconds": 0}, + CONF_KEEP_LAST_SAMPLE: False, + CONF_PERCENTILE: 50.0, + CONF_PRECISION: 2.0, + }, + title=DEFAULT_NAME, + ) + config_entry.add_to_hass(hass) + assert await hass.config_entries.async_setup(config_entry.entry_id) + await hass.async_block_till_done() + + result = await hass.config_entries.options.async_init(config_entry.entry_id) + assert result["type"] == FlowResultType.FORM + assert result["errors"] is None + assert result["preview"] == "statistics" + + await client.send_json_auto_id( + { + "type": "statistics/start_preview", + "flow_id": result["flow_id"], + "flow_type": "options_flow", + "user_input": { + CONF_SAMPLES_MAX_BUFFER_SIZE: 20.0, + CONF_MAX_AGE: {"hours": 8, "minutes": 0, "seconds": 0}, + CONF_KEEP_LAST_SAMPLE: False, + CONF_PERCENTILE: 50.0, + CONF_PRECISION: 2.0, + }, + } + ) + + msg = await client.receive_json() + assert msg["success"] + assert msg["result"] is None + + msg = await client.receive_json() + assert msg["event"] == snapshot + assert len(hass.states.async_all()) == 2 + + # add state for the tests + hass.states.async_set("sensor.test_monitored", "20") + await hass.async_block_till_done() + + msg = await client.receive_json() + assert msg["event"] == snapshot(name="updated") + + +async def test_options_flow_sensor_preview_config_entry_removed( + recorder_mock: Recorder, hass: HomeAssistant, hass_ws_client: WebSocketGenerator +) -> None: + """Test the option flow preview where the config entry is removed.""" + client = await hass_ws_client(hass) + + # Setup the config entry + config_entry = MockConfigEntry( + data={}, + domain=DOMAIN, + options={ + CONF_NAME: DEFAULT_NAME, + CONF_ENTITY_ID: "sensor.test_monitored", + CONF_STATE_CHARACTERISTIC: STAT_AVERAGE_LINEAR, + CONF_SAMPLES_MAX_BUFFER_SIZE: 20.0, + CONF_MAX_AGE: {"hours": 8, "minutes": 0, "seconds": 0}, + CONF_KEEP_LAST_SAMPLE: False, + CONF_PERCENTILE: 50.0, + CONF_PRECISION: 2.0, + }, + title=DEFAULT_NAME, + ) + config_entry.add_to_hass(hass) + assert await hass.config_entries.async_setup(config_entry.entry_id) + await hass.async_block_till_done() + + result = await hass.config_entries.options.async_init(config_entry.entry_id) + assert result["type"] == FlowResultType.FORM + assert result["errors"] is None + assert result["preview"] == "statistics" + + await hass.config_entries.async_remove(config_entry.entry_id) + + await client.send_json_auto_id( + { + "type": "statistics/start_preview", + "flow_id": result["flow_id"], + "flow_type": "options_flow", + "user_input": { + CONF_SAMPLES_MAX_BUFFER_SIZE: 25.0, + CONF_MAX_AGE: {"hours": 8, "minutes": 0, "seconds": 0}, + CONF_KEEP_LAST_SAMPLE: False, + CONF_PERCENTILE: 50.0, + CONF_PRECISION: 2.0, + }, + } + ) + msg = await client.receive_json() + assert not msg["success"] + assert msg["error"] == { + "code": "home_assistant_error", + "message": "Config entry not found", + }