Add preview to Threshold config & option flow (#117181)

Co-authored-by: Erik Montnemery <erik@montnemery.com>
Co-authored-by: Franck Nijhof <git@frenck.dev>

homeassistant/components/threshold/binary_sensor.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from collections.abc import Callable, Mapping
 import logging
 from typing import Any
 
@@ -22,7 +23,13 @@ from homeassistant.const import (
     STATE_UNAVAILABLE,
     STATE_UNKNOWN,
 )
-from homeassistant.core import Event, EventStateChangedData, HomeAssistant, callback
+from homeassistant.core import (
+    CALLBACK_TYPE,
+    Event,
+    EventStateChangedData,
+    HomeAssistant,
+    callback,
+)
 from homeassistant.helpers import (
     config_validation as cv,
     device_registry as dr,
@@ -111,7 +118,6 @@ async def async_setup_entry(
     async_add_entities(
         [
             ThresholdSensor(
-                hass,
                 entity_id,
                 name,
                 lower,
@@ -145,7 +151,7 @@ async def async_setup_platform(
     async_add_entities(
         [
             ThresholdSensor(
-                hass, entity_id, name, lower, upper, hysteresis, device_class, None
+                entity_id, name, lower, upper, hysteresis, device_class, None
             )
         ],
     )
@@ -167,7 +173,6 @@ class ThresholdSensor(BinarySensorEntity):
 
     def __init__(
         self,
-        hass: HomeAssistant,
         entity_id: str,
         name: str,
         lower: float | None,
@@ -178,6 +183,7 @@ class ThresholdSensor(BinarySensorEntity):
         device_info: DeviceInfo | None = None,
     ) -> None:
         """Initialize the Threshold sensor."""
+        self._preview_callback: Callable[[str, Mapping[str, Any]], None] | None = None
         self._attr_unique_id = unique_id
         self._attr_device_info = device_info
         self._entity_id = entity_id
@@ -193,9 +199,17 @@ class ThresholdSensor(BinarySensorEntity):
         self._state: bool | None = None
         self.sensor_value: float | None = None
 
+    async def async_added_to_hass(self) -> None:
+        """Run when entity about to be added to hass."""
+        self._async_setup_sensor()
+
+    @callback
+    def _async_setup_sensor(self) -> None:
+        """Set up the sensor and start tracking state changes."""
+
         def _update_sensor_state() -> None:
             """Handle sensor state changes."""
-            if (new_state := hass.states.get(self._entity_id)) is None:
+            if (new_state := self.hass.states.get(self._entity_id)) is None:
                 return
 
             try:
@@ -210,17 +224,26 @@ class ThresholdSensor(BinarySensorEntity):
 
             self._update_state()
 
+            if self._preview_callback:
+                calculated_state = self._async_calculate_state()
+                self._preview_callback(
+                    calculated_state.state, calculated_state.attributes
+                )
+
         @callback
         def async_threshold_sensor_state_listener(
             event: Event[EventStateChangedData],
         ) -> None:
             """Handle sensor state changes."""
             _update_sensor_state()
-            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()
 
         self.async_on_remove(
             async_track_state_change_event(
-                hass, [entity_id], async_threshold_sensor_state_listener
+                self.hass, [self._entity_id], async_threshold_sensor_state_listener
             )
         )
         _update_sensor_state()
@@ -305,3 +328,26 @@ class ThresholdSensor(BinarySensorEntity):
                 self._state_position = POSITION_IN_RANGE
                 self._state = True
             return
+
+    @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 if neither lower nor upper thresholds are set
+        if not self._entity_id or (
+            not hasattr(self, "_threshold_lower")
+            and not hasattr(self, "_threshold_upper")
+        ):
+            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_setup_sensor()
+        return self._call_on_remove_callbacks

homeassistant/components/threshold/config_flow.py

@@ -7,8 +7,11 @@ from typing import Any
 
 import voluptuous as vol
 
+from homeassistant.components import websocket_api
 from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN
 from homeassistant.const import CONF_ENTITY_ID, CONF_NAME
+from homeassistant.core import HomeAssistant, callback
+from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers import selector
 from homeassistant.helpers.schema_config_entry_flow import (
     SchemaCommonFlowHandler,
@@ -17,6 +20,7 @@ from homeassistant.helpers.schema_config_entry_flow import (
     SchemaFlowFormStep,
 )
 
+from .binary_sensor import ThresholdSensor
 from .const import CONF_HYSTERESIS, CONF_LOWER, CONF_UPPER, DEFAULT_HYSTERESIS, DOMAIN
 
 
@@ -61,11 +65,15 @@ CONFIG_SCHEMA = vol.Schema(
 ).extend(OPTIONS_SCHEMA.schema)
 
 CONFIG_FLOW = {
-    "user": SchemaFlowFormStep(CONFIG_SCHEMA, validate_user_input=_validate_mode)
+    "user": SchemaFlowFormStep(
+        CONFIG_SCHEMA, preview="threshold", validate_user_input=_validate_mode
+    )
 }
 
 OPTIONS_FLOW = {
-    "init": SchemaFlowFormStep(OPTIONS_SCHEMA, validate_user_input=_validate_mode)
+    "init": SchemaFlowFormStep(
+        OPTIONS_SCHEMA, preview="threshold", validate_user_input=_validate_mode
+    )
 }
 
 
@@ -79,3 +87,61 @@ class ConfigFlowHandler(SchemaConfigFlowHandler, domain=DOMAIN):
         """Return config entry title."""
         name: str = options[CONF_NAME]
         return 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"): "threshold/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":
+        entity_id = msg["user_input"][CONF_ENTITY_ID]
+        name = msg["user_input"][CONF_NAME]
+    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]
+
+    @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}
+            )
+        )
+
+    preview_entity = ThresholdSensor(
+        entity_id,
+        name,
+        msg["user_input"].get(CONF_LOWER),
+        msg["user_input"].get(CONF_UPPER),
+        msg["user_input"].get(CONF_HYSTERESIS),
+        None,
+        None,
+    )
+    preview_entity.hass = hass
+
+    connection.send_result(msg["id"])
+    connection.subscriptions[msg["id"]] = preview_entity.async_start_preview(
+        async_preview_updated
+    )

tests/components/threshold/snapshots/test_config_flow.ambr

@@ -0,0 +1,47 @@
+# serializer version: 1
+# name: test_config_flow_preview_success[missing_entity_id]
+  dict({
+    'attributes': dict({
+      'friendly_name': '',
+    }),
+    'state': 'unavailable',
+  })
+# ---
+# name: test_config_flow_preview_success[missing_upper_lower]
+  dict({
+    'attributes': dict({
+      'friendly_name': 'Test Sensor',
+    }),
+    'state': 'unavailable',
+  })
+# ---
+# name: test_config_flow_preview_success[success]
+  dict({
+    'attributes': dict({
+      'entity_id': 'sensor.test_monitored',
+      'friendly_name': 'Test Sensor',
+      'hysteresis': 0.0,
+      'lower': 20.0,
+      'position': 'below',
+      'sensor_value': 16.0,
+      'type': 'lower',
+      'upper': None,
+    }),
+    'state': 'on',
+  })
+# ---
+# name: test_options_flow_preview
+  dict({
+    'attributes': dict({
+      'entity_id': 'sensor.test_monitored',
+      'friendly_name': 'Test Sensor',
+      'hysteresis': 0.0,
+      'lower': 20.0,
+      'position': 'below',
+      'sensor_value': 16.0,
+      'type': 'lower',
+      'upper': None,
+    }),
+    'state': 'on',
+  })
+# ---

tests/components/threshold/test_config_flow.py

@@ -3,13 +3,16 @@
 from unittest.mock import patch
 
 import pytest
+from syrupy import SnapshotAssertion
 
 from homeassistant import config_entries
 from homeassistant.components.threshold.const import DOMAIN
+from homeassistant.const import ATTR_UNIT_OF_MEASUREMENT, UnitOfTemperature
 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_config_flow(hass: HomeAssistant) -> None:
@@ -162,3 +165,183 @@ async def test_options(hass: HomeAssistant) -> None:
     state = hass.states.get("binary_sensor.my_threshold")
     assert state.state == "off"
     assert state.attributes["type"] == "upper"
+
+
+@pytest.mark.parametrize(
+    "user_input",
+    [
+        (
+            {
+                "name": "Test Sensor",
+                "entity_id": "sensor.test_monitored",
+                "hysteresis": 0.0,
+                "lower": 20.0,
+            }
+        ),
+        (
+            {
+                "name": "Test Sensor",
+                "entity_id": "sensor.test_monitored",
+                "hysteresis": 0.0,
+            }
+        ),
+        (
+            {
+                "name": "",
+                "entity_id": "",
+                "hysteresis": 0.0,
+                "lower": 20.0,
+            }
+        ),
+    ],
+    ids=("success", "missing_upper_lower", "missing_entity_id"),
+)
+async def test_config_flow_preview_success(
+    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,
+        {ATTR_UNIT_OF_MEASUREMENT: UnitOfTemperature.CELSIUS},
+    )
+
+    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
+    assert result["preview"] == "threshold"
+
+    await client.send_json_auto_id(
+        {
+            "type": "threshold/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(
+    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,
+        {ATTR_UNIT_OF_MEASUREMENT: UnitOfTemperature.CELSIUS},
+    )
+
+    # Setup the config entry
+    config_entry = MockConfigEntry(
+        data={},
+        domain=DOMAIN,
+        options={
+            "entity_id": "sensor.test_monitored",
+            "hysteresis": 0.0,
+            "lower": 20.0,
+            "name": "Test Sensor",
+            "upper": None,
+        },
+        title="Test Sensor",
+    )
+    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"] == "threshold"
+
+    await client.send_json_auto_id(
+        {
+            "type": "threshold/start_preview",
+            "flow_id": result["flow_id"],
+            "flow_type": "options_flow",
+            "user_input": {
+                "name": "Test Sensor",
+                "entity_id": "sensor.test_monitored",
+                "hysteresis": 0.0,
+                "lower": 20.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
+
+
+async def test_options_flow_sensor_preview_config_entry_removed(
+    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={
+            "entity_id": "sensor.test_monitored",
+            "hysteresis": 0.0,
+            "lower": 20.0,
+            "name": "Test Sensor",
+            "upper": None,
+        },
+        title="Test Sensor",
+    )
+    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"] == "threshold"
+
+    await hass.config_entries.async_remove(config_entry.entry_id)
+
+    await client.send_json_auto_id(
+        {
+            "type": "threshold/start_preview",
+            "flow_id": result["flow_id"],
+            "flow_type": "options_flow",
+            "user_input": {
+                "name": "Test Sensor",
+                "entity_id": "sensor.test_monitored",
+                "hysteresis": 0.0,
+                "lower": 20.0,
+            },
+        }
+    )
+    msg = await client.receive_json()
+    assert not msg["success"]
+    assert msg["error"] == {
+        "code": "home_assistant_error",
+        "message": "Config entry not found",
+    }