Use a single WS command for group preview (#98903)

* Use a single WS command for group preview

* Fix tests

homeassistant/components/group/binary_sensor.py

@@ -1,6 +1,8 @@
 """Platform allowing several binary sensor to be grouped into one binary sensor."""
 from __future__ import annotations
 
+from typing import Any
+
 import voluptuous as vol
 
 from homeassistant.components.binary_sensor import (
@@ -85,6 +87,20 @@ async def async_setup_entry(
     )
 
 
+@callback
+def async_create_preview_binary_sensor(
+    name: str, validated_config: dict[str, Any]
+) -> BinarySensorGroup:
+    """Create a preview sensor."""
+    return BinarySensorGroup(
+        None,
+        name,
+        None,
+        validated_config[CONF_ENTITIES],
+        validated_config[CONF_ALL],
+    )
+
+
 class BinarySensorGroup(GroupEntity, BinarySensorEntity):
     """Representation of a BinarySensorGroup."""
 

homeassistant/components/group/config_flow.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 
 from collections.abc import Callable, Coroutine, Mapping
 from functools import partial
-from typing import Any, Literal, cast
+from typing import Any, cast
 
 import voluptuous as vol
 
@@ -21,10 +21,10 @@ from homeassistant.helpers.schema_config_entry_flow import (
     entity_selector_without_own_entities,
 )
 
-from . import DOMAIN
+from . import DOMAIN, GroupEntity
-from .binary_sensor import CONF_ALL, BinarySensorGroup
+from .binary_sensor import CONF_ALL, async_create_preview_binary_sensor
 from .const import CONF_HIDE_MEMBERS, CONF_IGNORE_NON_NUMERIC
-from .sensor import SensorGroup
+from .sensor import async_create_preview_sensor
 
 _STATISTIC_MEASURES = [
     "min",
@@ -171,8 +171,8 @@ CONFIG_FLOW = {
     "user": SchemaFlowMenuStep(GROUP_TYPES),
     "binary_sensor": SchemaFlowFormStep(
         BINARY_SENSOR_CONFIG_SCHEMA,
+        preview="group",
         validate_user_input=set_group_type("binary_sensor"),
-        preview="group_binary_sensor",
     ),
     "cover": SchemaFlowFormStep(
         basic_group_config_schema("cover"),
@@ -196,8 +196,8 @@ CONFIG_FLOW = {
     ),
     "sensor": SchemaFlowFormStep(
         SENSOR_CONFIG_SCHEMA,
+        preview="group",
         validate_user_input=set_group_type("sensor"),
-        preview="group_sensor",
     ),
     "switch": SchemaFlowFormStep(
         basic_group_config_schema("switch"),
@@ -210,22 +210,33 @@ OPTIONS_FLOW = {
     "init": SchemaFlowFormStep(next_step=choose_options_step),
     "binary_sensor": SchemaFlowFormStep(
         binary_sensor_options_schema,
-        preview="group_binary_sensor",
+        preview="group",
     ),
     "cover": SchemaFlowFormStep(partial(basic_group_options_schema, "cover")),
     "fan": SchemaFlowFormStep(partial(basic_group_options_schema, "fan")),
     "light": SchemaFlowFormStep(partial(light_switch_options_schema, "light")),
     "lock": SchemaFlowFormStep(partial(basic_group_options_schema, "lock")),
     "media_player": SchemaFlowFormStep(
-        partial(basic_group_options_schema, "media_player")
+        partial(basic_group_options_schema, "media_player"),
+        preview="group",
     ),
     "sensor": SchemaFlowFormStep(
         partial(sensor_options_schema, "sensor"),
-        preview="group_sensor",
+        preview="group",
     ),
     "switch": SchemaFlowFormStep(partial(light_switch_options_schema, "switch")),
 }
 
+PREVIEW_OPTIONS_SCHEMA: dict[str, vol.Schema] = {}
+
+CREATE_PREVIEW_ENTITY: dict[
+    str,
+    Callable[[str, dict[str, Any]], GroupEntity],
+] = {
+    "binary_sensor": async_create_preview_binary_sensor,
+    "sensor": async_create_preview_sensor,
+}
+
 
 class GroupConfigFlowHandler(SchemaConfigFlowHandler, domain=DOMAIN):
     """Handle a config or options flow for groups."""
@@ -261,12 +272,20 @@ class GroupConfigFlowHandler(SchemaConfigFlowHandler, domain=DOMAIN):
         )
         _async_hide_members(hass, options[CONF_ENTITIES], hidden_by)
 
-    @callback
     @staticmethod
-    def async_setup_preview(hass: HomeAssistant) -> None:
+    async def async_setup_preview(hass: HomeAssistant) -> None:
         """Set up preview WS API."""
-        websocket_api.async_register_command(hass, ws_preview_sensor)
+        for group_type, form_step in OPTIONS_FLOW.items():
-        websocket_api.async_register_command(hass, ws_preview_binary_sensor)
+            if group_type not in GROUP_TYPES:
+                continue
+            schema = cast(
+                Callable[
+                    [SchemaCommonFlowHandler | None], Coroutine[Any, Any, vol.Schema]
+                ],
+                form_step.schema,
+            )
+            PREVIEW_OPTIONS_SCHEMA[group_type] = await schema(None)
+        websocket_api.async_register_command(hass, ws_start_preview)
 
 
 def _async_hide_members(
@@ -282,127 +301,51 @@ def _async_hide_members(
         registry.async_update_entity(entity_id, hidden_by=hidden_by)
 
 
+@websocket_api.websocket_command(
+    {
+        vol.Required("type"): "group/start_preview",
+        vol.Required("flow_id"): str,
+        vol.Required("flow_type"): vol.Any("config_flow", "options_flow"),
+        vol.Required("user_input"): dict,
+    }
+)
 @callback
-def _async_handle_ws_preview(
+def ws_start_preview(
     hass: HomeAssistant,
     connection: websocket_api.ActiveConnection,
     msg: dict[str, Any],
-    config_schema: vol.Schema,
-    options_schema: vol.Schema,
-    create_preview_entity: Callable[
-        [Literal["config_flow", "options_flow"], str, dict[str, Any]],
-        BinarySensorGroup | SensorGroup,
-    ],
 ) -> None:
     """Generate a preview."""
     if msg["flow_type"] == "config_flow":
-        validated = config_schema(msg["user_input"])
+        flow_status = hass.config_entries.flow.async_get(msg["flow_id"])
+        group_type = flow_status["step_id"]
+        form_step = cast(SchemaFlowFormStep, CONFIG_FLOW[group_type])
+        schema = cast(vol.Schema, form_step.schema)
+        validated = schema(msg["user_input"])
         name = validated["name"]
     else:
-        validated = options_schema(msg["user_input"])
         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
+        group_type = config_entry.options["group_type"]
         name = config_entry.options["name"]
+        validated = PREVIEW_OPTIONS_SCHEMA[group_type](msg["user_input"])
 
     @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"], {"state": state, "attributes": attributes}
+                msg["id"],
+                {"attributes": attributes, "group_type": group_type, "state": state},
             )
         )
 
-    preview_entity = create_preview_entity(msg["flow_type"], name, validated)
+    preview_entity = CREATE_PREVIEW_ENTITY[group_type](name, validated)
     preview_entity.hass = hass
 
     connection.send_result(msg["id"])
     connection.subscriptions[msg["id"]] = preview_entity.async_start_preview(
         async_preview_updated
     )
-
-
-@websocket_api.websocket_command(
-    {
-        vol.Required("type"): "group/binary_sensor/start_preview",
-        vol.Required("flow_id"): str,
-        vol.Required("flow_type"): vol.Any("config_flow", "options_flow"),
-        vol.Required("user_input"): dict,
-    }
-)
-@websocket_api.async_response
-async def ws_preview_binary_sensor(
-    hass: HomeAssistant, connection: websocket_api.ActiveConnection, msg: dict[str, Any]
-) -> None:
-    """Generate a preview."""
-
-    def create_preview_binary_sensor(
-        flow_type: Literal["config_flow", "options_flow"],
-        name: str,
-        validated_config: dict[str, Any],
-    ) -> BinarySensorGroup:
-        """Create a preview sensor."""
-        return BinarySensorGroup(
-            None,
-            name,
-            None,
-            validated_config[CONF_ENTITIES],
-            validated_config[CONF_ALL],
-        )
-
-    _async_handle_ws_preview(
-        hass,
-        connection,
-        msg,
-        BINARY_SENSOR_CONFIG_SCHEMA,
-        await binary_sensor_options_schema(None),
-        create_preview_binary_sensor,
-    )
-
-
-@websocket_api.websocket_command(
-    {
-        vol.Required("type"): "group/sensor/start_preview",
-        vol.Required("flow_id"): str,
-        vol.Required("flow_type"): vol.Any("config_flow", "options_flow"),
-        vol.Required("user_input"): dict,
-    }
-)
-@websocket_api.async_response
-async def ws_preview_sensor(
-    hass: HomeAssistant, connection: websocket_api.ActiveConnection, msg: dict[str, Any]
-) -> None:
-    """Generate a preview."""
-
-    def create_preview_sensor(
-        flow_type: Literal["config_flow", "options_flow"],
-        name: str,
-        validated_config: dict[str, Any],
-    ) -> SensorGroup:
-        """Create a preview sensor."""
-        ignore_non_numeric = (
-            False
-            if flow_type == "config_flow"
-            else validated_config[CONF_IGNORE_NON_NUMERIC]
-        )
-        return SensorGroup(
-            None,
-            name,
-            validated_config[CONF_ENTITIES],
-            ignore_non_numeric,
-            validated_config[CONF_TYPE],
-            None,
-            None,
-            None,
-        )
-
-    _async_handle_ws_preview(
-        hass,
-        connection,
-        msg,
-        SENSOR_CONFIG_SCHEMA,
-        await sensor_options_schema("sensor", None),
-        create_preview_sensor,
-    )

homeassistant/components/group/sensor.py

@@ -136,6 +136,23 @@ async def async_setup_entry(
     )
 
 
+@callback
+def async_create_preview_sensor(
+    name: str, validated_config: dict[str, Any]
+) -> SensorGroup:
+    """Create a preview sensor."""
+    return SensorGroup(
+        None,
+        name,
+        validated_config[CONF_ENTITIES],
+        validated_config.get(CONF_IGNORE_NON_NUMERIC, False),
+        validated_config[CONF_TYPE],
+        None,
+        None,
+        None,
+    )
+
+
 def calc_min(
     sensor_values: list[tuple[str, float, State]]
 ) -> tuple[dict[str, str | None], float | None]:

homeassistant/config_entries.py

@@ -1864,7 +1864,7 @@ class OptionsFlowManager(data_entry_flow.FlowManager):
         await _load_integration(self.hass, entry.domain, {})
         if entry.domain not in self._preview:
             self._preview.add(entry.domain)
-            flow.async_setup_preview(self.hass)
+            await flow.async_setup_preview(self.hass)
 
 
 class OptionsFlow(data_entry_flow.FlowHandler):

homeassistant/data_entry_flow.py

@@ -439,7 +439,7 @@ class FlowManager(abc.ABC):
         """Set up preview for a flow handler."""
         if flow.handler not in self._preview:
             self._preview.add(flow.handler)
-            flow.async_setup_preview(self.hass)
+            await flow.async_setup_preview(self.hass)
 
 
 class FlowHandler:
@@ -649,9 +649,8 @@ class FlowHandler:
     def async_remove(self) -> None:
         """Notification that the flow has been removed."""
 
-    @callback
     @staticmethod
-    def async_setup_preview(hass: HomeAssistant) -> None:
+    async def async_setup_preview(hass: HomeAssistant) -> None:
         """Set up preview."""
 
 

homeassistant/helpers/schema_config_entry_flow.py

@@ -292,9 +292,8 @@ class SchemaConfigFlowHandler(config_entries.ConfigFlow, ABC):
         """Initialize config flow."""
         self._common_handler = SchemaCommonFlowHandler(self, self.config_flow, None)
 
-    @callback
     @staticmethod
-    def async_setup_preview(hass: HomeAssistant) -> None:
+    async def async_setup_preview(hass: HomeAssistant) -> None:
         """Set up preview."""
 
     @classmethod
@@ -369,7 +368,8 @@ class SchemaOptionsFlowHandler(config_entries.OptionsFlowWithConfigEntry):
         options_flow: Mapping[str, SchemaFlowStep],
         async_options_flow_finished: Callable[[HomeAssistant, Mapping[str, Any]], None]
         | None = None,
-        async_setup_preview: Callable[[HomeAssistant], None] | None = None,
+        async_setup_preview: Callable[[HomeAssistant], Coroutine[Any, Any, None]]
+        | None = None,
     ) -> None:
         """Initialize options flow.
 

tests/components/group/test_config_flow.py

@@ -490,11 +490,11 @@ async def test_config_flow_preview(
     assert result["type"] == FlowResultType.FORM
     assert result["step_id"] == domain
     assert result["errors"] is None
-    assert result["preview"] == f"group_{domain}"
+    assert result["preview"] == "group"
 
     await client.send_json_auto_id(
         {
-            "type": f"group/{domain}/start_preview",
+            "type": "group/start_preview",
             "flow_id": result["flow_id"],
             "flow_type": "config_flow",
             "user_input": {"name": "My group", "entities": input_entities}
@@ -508,6 +508,7 @@ async def test_config_flow_preview(
     msg = await client.receive_json()
     assert msg["event"] == {
         "attributes": {"friendly_name": "My group"} | extra_attributes[0],
+        "group_type": domain,
         "state": "unavailable",
     }
 
@@ -522,8 +523,10 @@ async def test_config_flow_preview(
         }
         | extra_attributes[0]
         | extra_attributes[1],
+        "group_type": domain,
         "state": group_state,
     }
+    assert len(hass.states.async_all()) == 2
 
 
 @pytest.mark.parametrize(
@@ -582,14 +585,14 @@ async def test_option_flow_preview(
     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"] == f"group_{domain}"
+    assert result["preview"] == "group"
 
     hass.states.async_set(input_entities[0], input_states[0])
     hass.states.async_set(input_entities[1], input_states[1])
 
     await client.send_json_auto_id(
         {
-            "type": f"group/{domain}/start_preview",
+            "type": "group/start_preview",
             "flow_id": result["flow_id"],
             "flow_type": "options_flow",
             "user_input": {"entities": input_entities} | extra_user_input,
@@ -603,8 +606,10 @@ async def test_option_flow_preview(
     assert msg["event"] == {
         "attributes": {"entity_id": input_entities, "friendly_name": "My group"}
         | extra_attributes,
+        "group_type": domain,
         "state": group_state,
     }
+    assert len(hass.states.async_all()) == 3
 
 
 async def test_option_flow_sensor_preview_config_entry_removed(
@@ -635,13 +640,13 @@ async def test_option_flow_sensor_preview_config_entry_removed(
     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"] == "group_sensor"
+    assert result["preview"] == "group"
 
     await hass.config_entries.async_remove(config_entry.entry_id)
 
     await client.send_json_auto_id(
         {
-            "type": "group/sensor/start_preview",
+            "type": "group/start_preview",
             "flow_id": result["flow_id"],
             "flow_type": "options_flow",
             "user_input": {

tests/test_config_entries.py

@@ -3962,9 +3962,8 @@ async def test_preview_supported(
             """Mock Reauth."""
             return self.async_show_form(step_id="next", preview="test")
 
-        @callback
         @staticmethod
-        def async_setup_preview(hass: HomeAssistant) -> None:
+        async def async_setup_preview(hass: HomeAssistant) -> None:
             """Set up preview."""
             preview_calls.append(None)