Skip to content

Commit 2abc197

Browse files
abmantisCopilotarturpragacz
authored
Add extract_from_target websocket command (home-assistant#150124)
Co-authored-by: Copilot <[email protected]> Co-authored-by: Artur Pragacz <[email protected]>
1 parent a3dec46 commit 2abc197

File tree

2 files changed

+302
-2
lines changed

2 files changed

+302
-2
lines changed

‎homeassistant/components/websocket_api/commands.py‎

Lines changed: 40 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,12 @@
3434
TemplateError,
3535
Unauthorized,
3636
)
37-
from homeassistant.helpers import config_validation as cv, entity, template
37+
from homeassistant.helpers import (
38+
config_validation as cv,
39+
entity,
40+
target as target_helpers,
41+
template,
42+
)
3843
from homeassistant.helpers.condition import (
3944
async_get_all_descriptions as async_get_all_condition_descriptions,
4045
async_subscribe_platform_events as async_subscribe_condition_platform_events,
@@ -96,6 +101,7 @@ def async_register_commands(
96101
async_reg(hass, handle_call_service)
97102
async_reg(hass, handle_entity_source)
98103
async_reg(hass, handle_execute_script)
104+
async_reg(hass, handle_extract_from_target)
99105
async_reg(hass, handle_fire_event)
100106
async_reg(hass, handle_get_config)
101107
async_reg(hass, handle_get_services)
@@ -838,6 +844,39 @@ def handle_entity_source(
838844
connection.send_result(msg["id"], _serialize_entity_sources(entity_sources))
839845

840846

847+
@callback
848+
@decorators.websocket_command(
849+
{
850+
vol.Required("type"): "extract_from_target",
851+
vol.Required("target"): cv.TARGET_FIELDS,
852+
vol.Optional("expand_group", default=False): bool,
853+
}
854+
)
855+
def handle_extract_from_target(
856+
hass: HomeAssistant, connection: ActiveConnection, msg: dict[str, Any]
857+
) -> None:
858+
"""Handle extract from target command."""
859+
860+
selector_data = target_helpers.TargetSelectorData(msg["target"])
861+
extracted = target_helpers.async_extract_referenced_entity_ids(
862+
hass, selector_data, expand_group=msg["expand_group"]
863+
)
864+
865+
extracted_dict = {
866+
"referenced_entities": extracted.referenced.union(
867+
extracted.indirectly_referenced
868+
),
869+
"referenced_devices": extracted.referenced_devices,
870+
"referenced_areas": extracted.referenced_areas,
871+
"missing_devices": extracted.missing_devices,
872+
"missing_areas": extracted.missing_areas,
873+
"missing_floors": extracted.missing_floors,
874+
"missing_labels": extracted.missing_labels,
875+
}
876+
877+
connection.send_result(msg["id"], extracted_dict)
878+
879+
841880
@decorators.websocket_command(
842881
{
843882
vol.Required("type"): "subscribe_trigger",

‎tests/components/websocket_api/test_commands.py‎

Lines changed: 262 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,12 @@
3232
from homeassistant.const import SIGNAL_BOOTSTRAP_INTEGRATIONS
3333
from homeassistant.core import Context, HomeAssistant, State, SupportsResponse, callback
3434
from homeassistant.exceptions import HomeAssistantError, ServiceValidationError
35-
from homeassistant.helpers import device_registry as dr
35+
from homeassistant.helpers import (
36+
area_registry as ar,
37+
device_registry as dr,
38+
entity_registry as er,
39+
label_registry as lr,
40+
)
3641
from homeassistant.helpers.dispatcher import async_dispatcher_send
3742
from homeassistant.helpers.event import async_track_state_change_event
3843
from homeassistant.loader import Integration, async_get_integration
@@ -108,6 +113,29 @@ def _apply_entities_changes(state_dict: dict, change_dict: dict) -> None:
108113
del state_dict[STATE_KEY_LONG_NAMES[key]][item]
109114

110115

116+
def _assert_extract_from_target_command_result(
117+
msg: dict[str, Any],
118+
entities: set[str] | None = None,
119+
devices: set[str] | None = None,
120+
areas: set[str] | None = None,
121+
missing_devices: set[str] | None = None,
122+
missing_areas: set[str] | None = None,
123+
missing_labels: set[str] | None = None,
124+
missing_floors: set[str] | None = None,
125+
) -> None:
126+
assert msg["type"] == const.TYPE_RESULT
127+
assert msg["success"]
128+
129+
result = msg["result"]
130+
assert set(result["referenced_entities"]) == (entities or set())
131+
assert set(result["referenced_devices"]) == (devices or set())
132+
assert set(result["referenced_areas"]) == (areas or set())
133+
assert set(result["missing_devices"]) == (missing_devices or set())
134+
assert set(result["missing_areas"]) == (missing_areas or set())
135+
assert set(result["missing_floors"]) == (missing_floors or set())
136+
assert set(result["missing_labels"]) == (missing_labels or set())
137+
138+
111139
async def test_fire_event(
112140
hass: HomeAssistant, websocket_client: MockHAClientWebSocket
113141
) -> None:
@@ -3223,3 +3251,236 @@ async def mock_setup(hass: HomeAssistant, _) -> bool:
32233251

32243252
# The component has been loaded
32253253
assert "test" in hass.config.components
3254+
3255+
3256+
async def test_extract_from_target(
3257+
hass: HomeAssistant,
3258+
websocket_client: MockHAClientWebSocket,
3259+
area_registry: ar.AreaRegistry,
3260+
device_registry: dr.DeviceRegistry,
3261+
entity_registry: er.EntityRegistry,
3262+
label_registry: lr.LabelRegistry,
3263+
) -> None:
3264+
"""Test extract_from_target command with mixed target types including entities, devices, areas, and labels."""
3265+
3266+
async def call_command(target: dict[str, str]) -> Any:
3267+
await websocket_client.send_json_auto_id(
3268+
{"type": "extract_from_target", "target": target}
3269+
)
3270+
return await websocket_client.receive_json()
3271+
3272+
config_entry = MockConfigEntry(domain="test")
3273+
config_entry.add_to_hass(hass)
3274+
3275+
device1 = device_registry.async_get_or_create(
3276+
config_entry_id=config_entry.entry_id,
3277+
identifiers={("test", "device1")},
3278+
)
3279+
3280+
device2 = device_registry.async_get_or_create(
3281+
config_entry_id=config_entry.entry_id,
3282+
identifiers={("test", "device2")},
3283+
)
3284+
3285+
area_device = device_registry.async_get_or_create(
3286+
config_entry_id=config_entry.entry_id,
3287+
identifiers={("test", "device3")},
3288+
)
3289+
3290+
label2_device = device_registry.async_get_or_create(
3291+
config_entry_id=config_entry.entry_id,
3292+
identifiers={("test", "device4")},
3293+
)
3294+
3295+
kitchen_area = area_registry.async_create("Kitchen")
3296+
living_room_area = area_registry.async_create("Living Room")
3297+
label_area = area_registry.async_create("Bathroom")
3298+
label1 = label_registry.async_create("Test Label 1")
3299+
label2 = label_registry.async_create("Test Label 2")
3300+
3301+
# Associate devices with areas and labels
3302+
device_registry.async_update_device(area_device.id, area_id=kitchen_area.id)
3303+
device_registry.async_update_device(label2_device.id, labels={label2.label_id})
3304+
area_registry.async_update(label_area.id, labels={label1.label_id})
3305+
3306+
# Setup entities with targets
3307+
device1_entity1 = entity_registry.async_get_or_create(
3308+
"light", "test", "unique1", device_id=device1.id
3309+
)
3310+
device1_entity2 = entity_registry.async_get_or_create(
3311+
"switch", "test", "unique2", device_id=device1.id
3312+
)
3313+
device2_entity = entity_registry.async_get_or_create(
3314+
"sensor", "test", "unique3", device_id=device2.id
3315+
)
3316+
area_device_entity = entity_registry.async_get_or_create(
3317+
"light", "test", "unique4", device_id=area_device.id
3318+
)
3319+
area_entity = entity_registry.async_get_or_create("switch", "test", "unique5")
3320+
label_device_entity = entity_registry.async_get_or_create(
3321+
"light", "test", "unique6", device_id=label2_device.id
3322+
)
3323+
label_entity = entity_registry.async_get_or_create("switch", "test", "unique7")
3324+
3325+
# Associate entities with areas and labels
3326+
entity_registry.async_update_entity(
3327+
area_entity.entity_id, area_id=living_room_area.id
3328+
)
3329+
entity_registry.async_update_entity(
3330+
label_entity.entity_id, labels={label1.label_id}
3331+
)
3332+
3333+
msg = await call_command({"entity_id": ["light.unknown_entity"]})
3334+
_assert_extract_from_target_command_result(msg, entities={"light.unknown_entity"})
3335+
3336+
msg = await call_command({"device_id": [device1.id, device2.id]})
3337+
_assert_extract_from_target_command_result(
3338+
msg,
3339+
entities={
3340+
device1_entity1.entity_id,
3341+
device1_entity2.entity_id,
3342+
device2_entity.entity_id,
3343+
},
3344+
devices={device1.id, device2.id},
3345+
)
3346+
3347+
msg = await call_command({"area_id": [kitchen_area.id, living_room_area.id]})
3348+
_assert_extract_from_target_command_result(
3349+
msg,
3350+
entities={area_device_entity.entity_id, area_entity.entity_id},
3351+
areas={kitchen_area.id, living_room_area.id},
3352+
devices={area_device.id},
3353+
)
3354+
3355+
msg = await call_command({"label_id": [label1.label_id, label2.label_id]})
3356+
_assert_extract_from_target_command_result(
3357+
msg,
3358+
entities={label_device_entity.entity_id, label_entity.entity_id},
3359+
devices={label2_device.id},
3360+
areas={label_area.id},
3361+
)
3362+
3363+
# Test multiple mixed targets
3364+
msg = await call_command(
3365+
{
3366+
"entity_id": ["light.direct"],
3367+
"device_id": [device1.id],
3368+
"area_id": [kitchen_area.id],
3369+
"label_id": [label1.label_id],
3370+
},
3371+
)
3372+
_assert_extract_from_target_command_result(
3373+
msg,
3374+
entities={
3375+
"light.direct",
3376+
device1_entity1.entity_id,
3377+
device1_entity2.entity_id,
3378+
area_device_entity.entity_id,
3379+
label_entity.entity_id,
3380+
},
3381+
devices={device1.id, area_device.id},
3382+
areas={kitchen_area.id, label_area.id},
3383+
)
3384+
3385+
3386+
async def test_extract_from_target_expand_group(
3387+
hass: HomeAssistant, websocket_client: MockHAClientWebSocket
3388+
) -> None:
3389+
"""Test extract_from_target command with expand_group parameter."""
3390+
await async_setup_component(
3391+
hass,
3392+
"group",
3393+
{
3394+
"group": {
3395+
"test_group": {
3396+
"name": "Test Group",
3397+
"entities": ["light.kitchen", "light.living_room"],
3398+
}
3399+
}
3400+
},
3401+
)
3402+
3403+
hass.states.async_set("light.kitchen", "on")
3404+
hass.states.async_set("light.living_room", "off")
3405+
3406+
# Test without expand_group (default False)
3407+
await websocket_client.send_json_auto_id(
3408+
{
3409+
"type": "extract_from_target",
3410+
"target": {"entity_id": ["group.test_group"]},
3411+
}
3412+
)
3413+
msg = await websocket_client.receive_json()
3414+
_assert_extract_from_target_command_result(msg, entities={"group.test_group"})
3415+
3416+
# Test with expand_group=True
3417+
await websocket_client.send_json_auto_id(
3418+
{
3419+
"type": "extract_from_target",
3420+
"target": {"entity_id": ["group.test_group"]},
3421+
"expand_group": True,
3422+
}
3423+
)
3424+
msg = await websocket_client.receive_json()
3425+
_assert_extract_from_target_command_result(
3426+
msg,
3427+
entities={"light.kitchen", "light.living_room"},
3428+
)
3429+
3430+
3431+
async def test_extract_from_target_missing_entities(
3432+
hass: HomeAssistant, websocket_client: MockHAClientWebSocket
3433+
) -> None:
3434+
"""Test extract_from_target command with missing device IDs, area IDs, etc."""
3435+
await websocket_client.send_json_auto_id(
3436+
{
3437+
"type": "extract_from_target",
3438+
"target": {
3439+
"device_id": ["non_existent_device"],
3440+
"area_id": ["non_existent_area"],
3441+
"label_id": ["non_existent_label"],
3442+
},
3443+
}
3444+
)
3445+
3446+
msg = await websocket_client.receive_json()
3447+
# Non-existent devices/areas are still referenced but reported as missing
3448+
_assert_extract_from_target_command_result(
3449+
msg,
3450+
devices={"non_existent_device"},
3451+
areas={"non_existent_area"},
3452+
missing_areas={"non_existent_area"},
3453+
missing_devices={"non_existent_device"},
3454+
missing_labels={"non_existent_label"},
3455+
)
3456+
3457+
3458+
async def test_extract_from_target_empty_target(
3459+
hass: HomeAssistant, websocket_client: MockHAClientWebSocket
3460+
) -> None:
3461+
"""Test extract_from_target command with empty target."""
3462+
await websocket_client.send_json_auto_id(
3463+
{
3464+
"type": "extract_from_target",
3465+
"target": {},
3466+
}
3467+
)
3468+
3469+
msg = await websocket_client.receive_json()
3470+
_assert_extract_from_target_command_result(msg)
3471+
3472+
3473+
async def test_extract_from_target_validation_error(
3474+
hass: HomeAssistant, websocket_client: MockHAClientWebSocket
3475+
) -> None:
3476+
"""Test extract_from_target command with invalid target data."""
3477+
await websocket_client.send_json_auto_id(
3478+
{
3479+
"type": "extract_from_target",
3480+
"target": "invalid", # Should be a dict, not string
3481+
}
3482+
)
3483+
msg = await websocket_client.receive_json()
3484+
assert msg["type"] == const.TYPE_RESULT
3485+
assert not msg["success"]
3486+
assert "error" in msg

0 commit comments

Comments
 (0)