From 8308be185e6bb9d244eed83e996a2778f10267c8 Mon Sep 17 00:00:00 2001 From: abmantis Date: Wed, 6 Aug 2025 15:50:16 +0100 Subject: [PATCH] Add expand_target websocket command --- .../components/websocket_api/commands.py | 42 ++- .../components/websocket_api/test_commands.py | 263 +++++++++++++++++- 2 files changed, 303 insertions(+), 2 deletions(-) diff --git a/homeassistant/components/websocket_api/commands.py b/homeassistant/components/websocket_api/commands.py index b63e5e14820..69c6389776b 100644 --- a/homeassistant/components/websocket_api/commands.py +++ b/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", diff --git a/tests/components/websocket_api/test_commands.py b/tests/components/websocket_api/test_commands.py index 263cd4a4ed8..3de50f55454 100644 --- a/tests/components/websocket_api/test_commands.py +++ b/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