Add preview to statistics (#122590)

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
+    )

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,6 +405,12 @@ class StatisticsSensor(SensorEntity):
             return
         self._add_state_to_queue(new_state)
         self._async_purge_update_and_schedule()
+
+        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
@@ -604,6 +633,8 @@ class StatisticsSensor(SensorEntity):
         _LOGGER.debug("%s: executing scheduled update", self.entity_id)
         self._async_cancel_update_listener()
         self._async_purge_update_and_schedule()
+        # 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]:
@@ -648,6 +679,12 @@ class StatisticsSensor(SensorEntity):
                 self._add_state_to_queue(state)
 
         self._async_purge_update_and_schedule()
+
+        # 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)
 

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',
+  })
+# ---

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",
+    }