Add expand_target websocket command

homeassistant/components/websocket_api/commands.py

@@ -34,7 +34,12 @@ from homeassistant.exceptions import (
     TemplateError,
     Unauthorized,
 )
-from homeassistant.helpers import config_validation as cv, entity, template
+from homeassistant.helpers import (
+    config_validation as cv,
+    entity,
+    target as target_helpers,
+    template,
+)
 from homeassistant.helpers.condition import (
     async_get_all_descriptions as async_get_all_condition_descriptions,
     async_subscribe_platform_events as async_subscribe_condition_platform_events,
@@ -96,6 +101,7 @@ def async_register_commands(
     async_reg(hass, handle_call_service)
     async_reg(hass, handle_entity_source)
     async_reg(hass, handle_execute_script)
+    async_reg(hass, handle_expand_target)
     async_reg(hass, handle_fire_event)
     async_reg(hass, handle_get_config)
     async_reg(hass, handle_get_services)
@@ -829,6 +835,40 @@ def handle_entity_source(
     connection.send_result(msg["id"], _serialize_entity_sources(entity_sources))
 
 
+@callback
+@decorators.websocket_command(
+    {
+        vol.Required("type"): "expand_target",
+        vol.Required("target"): cv.ENTITY_SERVICE_FIELDS,
+        vol.Optional("expand_group", default=False): bool,
+    }
+)
+def handle_expand_target(
+    hass: HomeAssistant, connection: ActiveConnection, msg: dict[str, Any]
+) -> None:
+    """Handle expand target command."""
+
+    selector_data = target_helpers.TargetSelectorData(msg["target"])
+    extracted = target_helpers.async_extract_referenced_entity_ids(
+        hass, selector_data, expand_group=msg["expand_group"]
+    )
+
+    extracted_dict = {
+        "referenced_entities": extracted.referenced.union(
+            extracted.indirectly_referenced
+        ),
+        "referenced_devices": extracted.referenced_devices,
+        "referenced_areas": extracted.referenced_areas,
+        "missing_devices": extracted.missing_devices,
+        "missing_areas": extracted.missing_areas,
+        "missing_floors": extracted.missing_floors,
+        "missing_labels": extracted.missing_labels,
+    }
+
+    payload = json_bytes(extracted_dict)
+    connection.send_message(construct_result_message(msg["id"], payload))
+
+
 @decorators.websocket_command(
     {
         vol.Required("type"): "subscribe_trigger",

tests/components/websocket_api/test_commands.py

@@ -28,7 +28,12 @@ from homeassistant.config_entries import ConfigEntryState
 from homeassistant.const import SIGNAL_BOOTSTRAP_INTEGRATIONS
 from homeassistant.core import Context, HomeAssistant, State, SupportsResponse, callback
 from homeassistant.exceptions import HomeAssistantError, ServiceValidationError
-from homeassistant.helpers import device_registry as dr
+from homeassistant.helpers import (
+    area_registry as ar,
+    device_registry as dr,
+    entity_registry as er,
+    label_registry as lr,
+)
 from homeassistant.helpers.dispatcher import async_dispatcher_send
 from homeassistant.helpers.event import async_track_state_change_event
 from homeassistant.loader import Integration, async_get_integration
@@ -104,6 +109,29 @@ def _apply_entities_changes(state_dict: dict, change_dict: dict) -> None:
             del state_dict[STATE_KEY_LONG_NAMES[key]][item]
 
 
+def _assert_expand_target_command_result(
+    msg: dict[str, Any],
+    entities: set[str] | None = None,
+    devices: set[str] | None = None,
+    areas: set[str] | None = None,
+    missing_devices: set[str] | None = None,
+    missing_areas: set[str] | None = None,
+    missing_labels: set[str] | None = None,
+    missing_floors: set[str] | None = None,
+) -> None:
+    assert msg["type"] == const.TYPE_RESULT
+    assert msg["success"]
+
+    result = msg["result"]
+    assert set(result["referenced_entities"]) == (entities or set())
+    assert set(result["referenced_devices"]) == (devices or set())
+    assert set(result["referenced_areas"]) == (areas or set())
+    assert set(result["missing_devices"]) == (missing_devices or set())
+    assert set(result["missing_areas"]) == (missing_areas or set())
+    assert set(result["missing_floors"]) == (missing_floors or set())
+    assert set(result["missing_labels"]) == (missing_labels or set())
+
+
 async def test_fire_event(
     hass: HomeAssistant, websocket_client: MockHAClientWebSocket
 ) -> None:
@@ -3115,3 +3143,236 @@ async def test_wait_integration_startup(
 
     # The component has been loaded
     assert "test" in hass.config.components
+
+
+async def test_expand_target(
+    hass: HomeAssistant,
+    websocket_client: MockHAClientWebSocket,
+    area_registry: ar.AreaRegistry,
+    device_registry: dr.DeviceRegistry,
+    entity_registry: er.EntityRegistry,
+    label_registry: lr.LabelRegistry,
+) -> None:
+    """Test expand_target command with mixed target types including entities, devices, areas, and labels."""
+
+    async def call_command(target: dict[str, str]) -> Any:
+        await websocket_client.send_json_auto_id(
+            {"type": "expand_target", "target": target}
+        )
+        return await websocket_client.receive_json()
+
+    config_entry = MockConfigEntry(domain="test")
+    config_entry.add_to_hass(hass)
+
+    device1 = device_registry.async_get_or_create(
+        config_entry_id=config_entry.entry_id,
+        identifiers={("test", "device1")},
+    )
+
+    device2 = device_registry.async_get_or_create(
+        config_entry_id=config_entry.entry_id,
+        identifiers={("test", "device2")},
+    )
+
+    area_device = device_registry.async_get_or_create(
+        config_entry_id=config_entry.entry_id,
+        identifiers={("test", "device3")},
+    )
+
+    label2_device = device_registry.async_get_or_create(
+        config_entry_id=config_entry.entry_id,
+        identifiers={("test", "device4")},
+    )
+
+    kitchen_area = area_registry.async_create("Kitchen")
+    living_room_area = area_registry.async_create("Living Room")
+    label_area = area_registry.async_create("Bathroom")
+    label1 = label_registry.async_create("Test Label 1")
+    label2 = label_registry.async_create("Test Label 2")
+
+    # Associate devices with areas and labels
+    device_registry.async_update_device(area_device.id, area_id=kitchen_area.id)
+    device_registry.async_update_device(label2_device.id, labels={label2.label_id})
+    area_registry.async_update(label_area.id, labels={label1.label_id})
+
+    # Setup entities with targets
+    device1_entity1 = entity_registry.async_get_or_create(
+        "light", "test", "unique1", device_id=device1.id
+    )
+    device1_entity2 = entity_registry.async_get_or_create(
+        "switch", "test", "unique2", device_id=device1.id
+    )
+    device2_entity = entity_registry.async_get_or_create(
+        "sensor", "test", "unique3", device_id=device2.id
+    )
+    area_device_entity = entity_registry.async_get_or_create(
+        "light", "test", "unique4", device_id=area_device.id
+    )
+    area_entity = entity_registry.async_get_or_create("switch", "test", "unique5")
+    label_device_entity = entity_registry.async_get_or_create(
+        "light", "test", "unique6", device_id=label2_device.id
+    )
+    label_entity = entity_registry.async_get_or_create("switch", "test", "unique7")
+
+    # Associate entities with areas and labels
+    entity_registry.async_update_entity(
+        area_entity.entity_id, area_id=living_room_area.id
+    )
+    entity_registry.async_update_entity(
+        label_entity.entity_id, labels={label1.label_id}
+    )
+
+    msg = await call_command({"entity_id": ["light.unknown_entity"]})
+    _assert_expand_target_command_result(msg, entities={"light.unknown_entity"})
+
+    msg = await call_command({"device_id": [device1.id, device2.id]})
+    _assert_expand_target_command_result(
+        msg,
+        entities={
+            device1_entity1.entity_id,
+            device1_entity2.entity_id,
+            device2_entity.entity_id,
+        },
+        devices={device1.id, device2.id},
+    )
+
+    msg = await call_command({"area_id": [kitchen_area.id, living_room_area.id]})
+    _assert_expand_target_command_result(
+        msg,
+        entities={area_device_entity.entity_id, area_entity.entity_id},
+        areas={kitchen_area.id, living_room_area.id},
+        devices={area_device.id},
+    )
+
+    msg = await call_command({"label_id": [label1.label_id, label2.label_id]})
+    _assert_expand_target_command_result(
+        msg,
+        entities={label_device_entity.entity_id, label_entity.entity_id},
+        devices={label2_device.id},
+        areas={label_area.id},
+    )
+
+    # Test multiple mixed targets
+    msg = await call_command(
+        {
+            "entity_id": ["light.direct"],
+            "device_id": [device1.id],
+            "area_id": [kitchen_area.id],
+            "label_id": [label1.label_id],
+        },
+    )
+    _assert_expand_target_command_result(
+        msg,
+        entities={
+            "light.direct",
+            device1_entity1.entity_id,
+            device1_entity2.entity_id,
+            area_device_entity.entity_id,
+            label_entity.entity_id,
+        },
+        devices={device1.id, area_device.id},
+        areas={kitchen_area.id, label_area.id},
+    )
+
+
+async def test_expand_target_expand_group(
+    hass: HomeAssistant, websocket_client: MockHAClientWebSocket
+) -> None:
+    """Test expand_target command with expand_group parameter."""
+    await async_setup_component(
+        hass,
+        "group",
+        {
+            "group": {
+                "test_group": {
+                    "name": "Test Group",
+                    "entities": ["light.kitchen", "light.living_room"],
+                }
+            }
+        },
+    )
+
+    hass.states.async_set("light.kitchen", "on")
+    hass.states.async_set("light.living_room", "off")
+
+    # Test without expand_group (default False)
+    await websocket_client.send_json_auto_id(
+        {
+            "type": "expand_target",
+            "target": {"entity_id": ["group.test_group"]},
+        }
+    )
+    msg = await websocket_client.receive_json()
+    _assert_expand_target_command_result(msg, entities={"group.test_group"})
+
+    # Test with expand_group=True
+    await websocket_client.send_json_auto_id(
+        {
+            "type": "expand_target",
+            "target": {"entity_id": ["group.test_group"]},
+            "expand_group": True,
+        }
+    )
+    msg = await websocket_client.receive_json()
+    _assert_expand_target_command_result(
+        msg,
+        entities={"light.kitchen", "light.living_room"},
+    )
+
+
+async def test_expand_target_missing_entities(
+    hass: HomeAssistant, websocket_client: MockHAClientWebSocket
+) -> None:
+    """Test expand_target command with missing device IDs, area IDs, etc."""
+    await websocket_client.send_json_auto_id(
+        {
+            "type": "expand_target",
+            "target": {
+                "device_id": ["non_existent_device"],
+                "area_id": ["non_existent_area"],
+                "label_id": ["non_existent_label"],
+            },
+        }
+    )
+
+    msg = await websocket_client.receive_json()
+    # Non-existent devices/areas are still referenced but reported as missing
+    _assert_expand_target_command_result(
+        msg,
+        devices={"non_existent_device"},
+        areas={"non_existent_area"},
+        missing_areas={"non_existent_area"},
+        missing_devices={"non_existent_device"},
+        missing_labels={"non_existent_label"},
+    )
+
+
+async def test_expand_target_empty_target(
+    hass: HomeAssistant, websocket_client: MockHAClientWebSocket
+) -> None:
+    """Test expand_target command with empty target."""
+    await websocket_client.send_json_auto_id(
+        {
+            "type": "expand_target",
+            "target": {},
+        }
+    )
+
+    msg = await websocket_client.receive_json()
+    _assert_expand_target_command_result(msg)
+
+
+async def test_expand_target_validation_error(
+    hass: HomeAssistant, websocket_client: MockHAClientWebSocket
+) -> None:
+    """Test expand_target command with invalid target data."""
+    await websocket_client.send_json_auto_id(
+        {
+            "type": "expand_target",
+            "target": "invalid",  # Should be a dict, not string
+        }
+    )
+    msg = await websocket_client.receive_json()
+    assert msg["type"] == const.TYPE_RESULT
+    assert not msg["success"]
+    assert "error" in msg