From a5e5c5c7e6407ff0d7c9756a92c31b9776168712 Mon Sep 17 00:00:00 2001 From: Erik Montnemery Date: Fri, 27 Nov 2020 09:03:44 +0100 Subject: [PATCH] Exclude disabled entities from async_entries_for_device (#43665) --- homeassistant/components/deconz/services.py | 9 ++++- homeassistant/components/mqtt/__init__.py | 2 +- homeassistant/components/mqtt/debug_info.py | 2 +- homeassistant/components/tasmota/discovery.py | 4 ++- homeassistant/components/tuya/__init__.py | 2 +- .../components/unifi/unifi_entity_base.py | 11 +++++- .../components/zha/core/discovery.py | 4 ++- homeassistant/components/zha/core/gateway.py | 4 ++- homeassistant/components/zha/core/group.py | 4 ++- homeassistant/helpers/entity_registry.py | 15 +++++--- homeassistant/helpers/service.py | 2 +- tests/helpers/test_entity_registry.py | 36 +++++++++++++++++++ 12 files changed, 81 insertions(+), 14 deletions(-) diff --git a/homeassistant/components/deconz/services.py b/homeassistant/components/deconz/services.py index a90f770eb9b..65b9c1aad8a 100644 --- a/homeassistant/components/deconz/services.py +++ b/homeassistant/components/deconz/services.py @@ -207,5 +207,12 @@ async def async_remove_orphaned_entries_service(hass, data): # Remove devices that don't belong to any entity for device_id in devices_to_be_removed: - if len(async_entries_for_device(entity_registry, device_id)) == 0: + if ( + len( + async_entries_for_device( + entity_registry, device_id, include_disabled_entities=True + ) + ) + == 0 + ): device_registry.async_remove_device(device_id) diff --git a/homeassistant/components/mqtt/__init__.py b/homeassistant/components/mqtt/__init__.py index 27b43142f7c..bd89a2364be 100644 --- a/homeassistant/components/mqtt/__init__.py +++ b/homeassistant/components/mqtt/__init__.py @@ -1248,7 +1248,7 @@ async def cleanup_device_registry(hass, device_id): if ( device_id and not hass.helpers.entity_registry.async_entries_for_device( - entity_registry, device_id + entity_registry, device_id, include_disabled_entities=True ) and not await device_trigger.async_get_triggers(hass, device_id) and not tag.async_has_tags(hass, device_id) diff --git a/homeassistant/components/mqtt/debug_info.py b/homeassistant/components/mqtt/debug_info.py index 75e4b53a191..a3c56652253 100644 --- a/homeassistant/components/mqtt/debug_info.py +++ b/homeassistant/components/mqtt/debug_info.py @@ -123,7 +123,7 @@ async def info_for_device(hass, device_id): entity_registry = await hass.helpers.entity_registry.async_get_registry() entries = hass.helpers.entity_registry.async_entries_for_device( - entity_registry, device_id + entity_registry, device_id, include_disabled_entities=True ) mqtt_debug_info = hass.data.setdefault( DATA_MQTT_DEBUG_INFO, {"entities": {}, "triggers": {}} diff --git a/homeassistant/components/tasmota/discovery.py b/homeassistant/components/tasmota/discovery.py index 2313a8327c5..22824e9cd71 100644 --- a/homeassistant/components/tasmota/discovery.py +++ b/homeassistant/components/tasmota/discovery.py @@ -144,7 +144,9 @@ async def async_start( orphaned_entities = { entry.unique_id - for entry in async_entries_for_device(entity_registry, device.id) + for entry in async_entries_for_device( + entity_registry, device.id, include_disabled_entities=True + ) if entry.domain == sensor.DOMAIN and entry.platform == DOMAIN } for (tasmota_sensor_config, discovery_hash) in sensors: diff --git a/homeassistant/components/tuya/__init__.py b/homeassistant/components/tuya/__init__.py index 64a2d203695..bc665baeb86 100644 --- a/homeassistant/components/tuya/__init__.py +++ b/homeassistant/components/tuya/__init__.py @@ -269,7 +269,7 @@ async def cleanup_device_registry(hass: HomeAssistant, device_id): device_registry = await hass.helpers.device_registry.async_get_registry() entity_registry = await hass.helpers.entity_registry.async_get_registry() if device_id and not hass.helpers.entity_registry.async_entries_for_device( - entity_registry, device_id + entity_registry, device_id, include_disabled_entities=True ): device_registry.async_remove_device(device_id) diff --git a/homeassistant/components/unifi/unifi_entity_base.py b/homeassistant/components/unifi/unifi_entity_base.py index a730c134603..7b45d309c14 100644 --- a/homeassistant/components/unifi/unifi_entity_base.py +++ b/homeassistant/components/unifi/unifi_entity_base.py @@ -95,7 +95,16 @@ class UniFiBase(Entity): entity_registry.async_remove(self.entity_id) return - if len(async_entries_for_device(entity_registry, entity_entry.device_id)) == 1: + if ( + len( + async_entries_for_device( + entity_registry, + entity_entry.device_id, + include_disabled_entities=True, + ) + ) + == 1 + ): device_registry.async_remove_device(device_entry.id) return diff --git a/homeassistant/components/zha/core/discovery.py b/homeassistant/components/zha/core/discovery.py index 4dff2c6b16b..05a12bc2284 100644 --- a/homeassistant/components/zha/core/discovery.py +++ b/homeassistant/components/zha/core/discovery.py @@ -243,7 +243,9 @@ class GroupProbe: if member.device.is_coordinator: continue entities = async_entries_for_device( - zha_gateway.ha_entity_registry, member.device.device_id + zha_gateway.ha_entity_registry, + member.device.device_id, + include_disabled_entities=True, ) all_domain_occurrences.extend( [ diff --git a/homeassistant/components/zha/core/gateway.py b/homeassistant/components/zha/core/gateway.py index 69bf85e47a4..13a87c13226 100644 --- a/homeassistant/components/zha/core/gateway.py +++ b/homeassistant/components/zha/core/gateway.py @@ -400,7 +400,9 @@ class ZHAGateway: # then we get all group entity entries tied to the coordinator all_group_entity_entries = async_entries_for_device( - self.ha_entity_registry, self.coordinator_zha_device.device_id + self.ha_entity_registry, + self.coordinator_zha_device.device_id, + include_disabled_entities=True, ) # then we get the entity entries for this specific group by getting the entries that match diff --git a/homeassistant/components/zha/core/group.py b/homeassistant/components/zha/core/group.py index 2961f335989..8edb1da8f68 100644 --- a/homeassistant/components/zha/core/group.py +++ b/homeassistant/components/zha/core/group.py @@ -195,7 +195,9 @@ class ZHAGroup(LogMixin): domain_entity_ids: List[str] = [] for member in self.members: entities = async_entries_for_device( - self._zha_gateway.ha_entity_registry, member.device.device_id + self._zha_gateway.ha_entity_registry, + member.device.device_id, + include_disabled_entities=True, ) domain_entity_ids.extend( [entity.entity_id for entity in entities if entity.domain == domain] diff --git a/homeassistant/helpers/entity_registry.py b/homeassistant/helpers/entity_registry.py index 143f3a99137..7e8700e8236 100644 --- a/homeassistant/helpers/entity_registry.py +++ b/homeassistant/helpers/entity_registry.py @@ -298,7 +298,9 @@ class EntityRegistry: the device is disabled. """ if event.data["action"] == "remove": - entities = async_entries_for_device(self, event.data["device_id"]) + entities = async_entries_for_device( + self, event.data["device_id"], include_disabled_entities=True + ) for entity in entities: self.async_remove(entity.entity_id) return @@ -311,7 +313,9 @@ class EntityRegistry: if not device.disabled: return - entities = async_entries_for_device(self, event.data["device_id"]) + entities = async_entries_for_device( + self, event.data["device_id"], include_disabled_entities=True + ) for entity in entities: self.async_update_entity( # type: ignore entity.entity_id, disabled_by=DISABLED_DEVICE @@ -548,11 +552,14 @@ async def async_get_registry(hass: HomeAssistantType) -> EntityRegistry: @callback def async_entries_for_device( - registry: EntityRegistry, device_id: str + registry: EntityRegistry, device_id: str, include_disabled_entities: bool = False ) -> List[RegistryEntry]: """Return entries that match a device.""" return [ - entry for entry in registry.entities.values() if entry.device_id == device_id + entry + for entry in registry.entities.values() + if entry.device_id == device_id + and (not entry.disabled_by or include_disabled_entities) ] diff --git a/homeassistant/helpers/service.py b/homeassistant/helpers/service.py index 6a290a77b08..e6c805e9a35 100644 --- a/homeassistant/helpers/service.py +++ b/homeassistant/helpers/service.py @@ -260,7 +260,7 @@ async def async_extract_entity_ids( entry.entity_id for device in devices for entry in hass.helpers.entity_registry.async_entries_for_device( - ent_reg, device.id + ent_reg, device.id, include_disabled_entities=True ) if not entry.area_id ) diff --git a/tests/helpers/test_entity_registry.py b/tests/helpers/test_entity_registry.py index f42661ec915..960537e784c 100644 --- a/tests/helpers/test_entity_registry.py +++ b/tests/helpers/test_entity_registry.py @@ -736,3 +736,39 @@ async def test_disable_device_disables_entities(hass, registry): entry = registry.async_get(entry.entity_id) assert entry.disabled assert entry.disabled_by == "device" + + +async def test_disabled_entities_excluded_from_entity_list(hass, registry): + """Test that disabled entities are exclduded from async_entries_for_device.""" + device_registry = mock_device_registry(hass) + config_entry = MockConfigEntry(domain="light") + + device_entry = device_registry.async_get_or_create( + config_entry_id=config_entry.entry_id, + connections={("mac", "12:34:56:AB:CD:EF")}, + ) + + entry1 = registry.async_get_or_create( + "light", + "hue", + "5678", + config_entry=config_entry, + device_id=device_entry.id, + ) + + entry2 = registry.async_get_or_create( + "light", + "hue", + "ABCD", + config_entry=config_entry, + device_id=device_entry.id, + disabled_by="user", + ) + + entries = entity_registry.async_entries_for_device(registry, device_entry.id) + assert entries == [entry1] + + entries = entity_registry.async_entries_for_device( + registry, device_entry.id, include_disabled_entities=True + ) + assert entries == [entry1, entry2]