From 330a7afdfc726f44b8f75cb4da6d06d467bce866 Mon Sep 17 00:00:00 2001 From: Erik Montnemery Date: Wed, 26 Apr 2023 18:42:49 +0200 Subject: [PATCH] Teach switch_as_x about exposed entities (#92059) --- .../homeassistant/exposed_entities.py | 36 ++++ .../components/switch_as_x/__init__.py | 35 +++- .../components/switch_as_x/entity.py | 34 +++- tests/components/switch_as_x/test_init.py | 192 +++++++++++++++++- 4 files changed, 285 insertions(+), 12 deletions(-) diff --git a/homeassistant/components/homeassistant/exposed_entities.py b/homeassistant/components/homeassistant/exposed_entities.py index 852e33721c0..b246c26e91c 100644 --- a/homeassistant/components/homeassistant/exposed_entities.py +++ b/homeassistant/components/homeassistant/exposed_entities.py @@ -156,6 +156,21 @@ class ExposedEntities: return result + @callback + def async_get_entity_settings(self, entity_id: str) -> dict[str, Mapping[str, Any]]: + """Get assistant expose settings for an entity.""" + entity_registry = er.async_get(self._hass) + result: dict[str, Mapping[str, Any]] = {} + + if not (registry_entry := entity_registry.async_get(entity_id)): + raise HomeAssistantError("Unknown entity") + + for assistant in KNOWN_ASSISTANTS: + if options := registry_entry.options.get(assistant): + result[assistant] = options + + return result + @callback def async_should_expose(self, assistant: str, entity_id: str) -> bool: """Return True if an entity should be exposed to an assistant.""" @@ -348,6 +363,27 @@ def async_get_assistant_settings( return exposed_entities.async_get_assistant_settings(assistant) +@callback +def async_get_entity_settings( + hass: HomeAssistant, entity_id: str +) -> dict[str, Mapping[str, Any]]: + """Get assistant expose settings for an entity.""" + exposed_entities: ExposedEntities = hass.data[DATA_EXPOSED_ENTITIES] + return exposed_entities.async_get_entity_settings(entity_id) + + +@callback +def async_expose_entity( + hass: HomeAssistant, + assistant: str, + entity_id: str, + should_expose: bool, +) -> None: + """Get assistant expose settings for an entity.""" + exposed_entities: ExposedEntities = hass.data[DATA_EXPOSED_ENTITIES] + exposed_entities.async_expose_entity(assistant, entity_id, should_expose) + + @callback def async_should_expose(hass: HomeAssistant, assistant: str, entity_id: str) -> bool: """Return True if an entity should be exposed to an assistant.""" diff --git a/homeassistant/components/switch_as_x/__init__.py b/homeassistant/components/switch_as_x/__init__.py index 102319cec93..ef64a86c6e8 100644 --- a/homeassistant/components/switch_as_x/__init__.py +++ b/homeassistant/components/switch_as_x/__init__.py @@ -5,6 +5,7 @@ import logging import voluptuous as vol +from homeassistant.components.homeassistant import exposed_entities from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_ENTITY_ID from homeassistant.core import Event, HomeAssistant, callback @@ -104,17 +105,39 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: async def async_remove_entry(hass: HomeAssistant, entry: ConfigEntry) -> None: - """Unload a config entry.""" - # Unhide the wrapped entry if registered + """Unload a config entry. + + This will unhide the wrapped entity and restore assistant expose settings. + """ registry = er.async_get(hass) try: - entity_id = er.async_validate_entity_id(registry, entry.options[CONF_ENTITY_ID]) + switch_entity_id = er.async_validate_entity_id( + registry, entry.options[CONF_ENTITY_ID] + ) except vol.Invalid: # The source entity has been removed from the entity registry return - if not (entity_entry := registry.async_get(entity_id)): + if not (switch_entity_entry := registry.async_get(switch_entity_id)): return - if entity_entry.hidden_by == er.RegistryEntryHider.INTEGRATION: - registry.async_update_entity(entity_id, hidden_by=None) + # Unhide the wrapped entity + if switch_entity_entry.hidden_by == er.RegistryEntryHider.INTEGRATION: + registry.async_update_entity(switch_entity_id, hidden_by=None) + + switch_as_x_entries = er.async_entries_for_config_entry(registry, entry.entry_id) + if not switch_as_x_entries: + return + + switch_as_x_entry = switch_as_x_entries[0] + + # Restore assistant expose settings + expose_settings = exposed_entities.async_get_entity_settings( + hass, switch_as_x_entry.entity_id + ) + for assistant, settings in expose_settings.items(): + if (should_expose := settings.get("should_expose")) is None: + continue + exposed_entities.async_expose_entity( + hass, assistant, switch_entity_id, should_expose + ) diff --git a/homeassistant/components/switch_as_x/entity.py b/homeassistant/components/switch_as_x/entity.py index 21a7b882442..a73271bdc83 100644 --- a/homeassistant/components/switch_as_x/entity.py +++ b/homeassistant/components/switch_as_x/entity.py @@ -3,6 +3,7 @@ from __future__ import annotations from typing import Any +from homeassistant.components.homeassistant import exposed_entities from homeassistant.components.switch import DOMAIN as SWITCH_DOMAIN from homeassistant.const import ( ATTR_ENTITY_ID, @@ -99,14 +100,37 @@ class BaseEntity(Entity): {"entity_id": self._switch_entity_id}, ) - if not self._is_new_entity: + if not self._is_new_entity or not ( + wrapped_switch := registry.async_get(self._switch_entity_id) + ): return - wrapped_switch = registry.async_get(self._switch_entity_id) - if not wrapped_switch or wrapped_switch.name is None: - return + def copy_custom_name(wrapped_switch: er.RegistryEntry) -> None: + """Copy the name set by user from the wrapped entity.""" + if wrapped_switch.name is None: + return + registry.async_update_entity(self.entity_id, name=wrapped_switch.name) - registry.async_update_entity(self.entity_id, name=wrapped_switch.name) + def copy_expose_settings() -> None: + """Copy assistant expose settings from the wrapped entity. + + Also unexpose the wrapped entity if exposed. + """ + expose_settings = exposed_entities.async_get_entity_settings( + self.hass, self._switch_entity_id + ) + for assistant, settings in expose_settings.items(): + if (should_expose := settings.get("should_expose")) is None: + continue + exposed_entities.async_expose_entity( + self.hass, assistant, self.entity_id, should_expose + ) + exposed_entities.async_expose_entity( + self.hass, assistant, self._switch_entity_id, False + ) + + copy_custom_name(wrapped_switch) + copy_expose_settings() class BaseToggleEntity(BaseEntity, ToggleEntity): diff --git a/tests/components/switch_as_x/test_init.py b/tests/components/switch_as_x/test_init.py index 87cc291a599..fac744d0c0e 100644 --- a/tests/components/switch_as_x/test_init.py +++ b/tests/components/switch_as_x/test_init.py @@ -5,6 +5,7 @@ from unittest.mock import patch import pytest +from homeassistant.components.homeassistant import exposed_entities from homeassistant.components.switch_as_x.const import CONF_TARGET_DOMAIN, DOMAIN from homeassistant.const import ( CONF_ENTITY_ID, @@ -19,9 +20,16 @@ from homeassistant.const import ( ) from homeassistant.core import HomeAssistant from homeassistant.helpers import device_registry as dr, entity_registry as er +from homeassistant.setup import async_setup_component from tests.common import MockConfigEntry +EXPOSE_SETTINGS = { + "cloud.alexa": True, + "cloud.google_assistant": False, + "conversation": True, +} + PLATFORMS_TO_TEST = ( Platform.COVER, Platform.FAN, @@ -607,7 +615,7 @@ async def test_custom_name_2( ) -> None: """Test the source entity has a custom name. - This tests the custom name is only copied from the source device when the config + This tests the custom name is only copied from the source device when the switch_as_x config entry is setup the first time. """ registry = er.async_get(hass) @@ -647,6 +655,8 @@ async def test_custom_name_2( ) switch_as_x_config_entry.add_to_hass(hass) + # Register the switch as x entity in the entity registry, this means + # the entity has been setup before switch_as_x_entity_entry = registry.async_get_or_create( target_domain, "switch_as_x", @@ -674,3 +684,183 @@ async def test_custom_name_2( assert entity_entry.options == { DOMAIN: {"entity_id": switch_entity_entry.entity_id} } + + +@pytest.mark.parametrize("target_domain", PLATFORMS_TO_TEST) +async def test_import_expose_settings_1( + hass: HomeAssistant, + target_domain: Platform, +) -> None: + """Test importing assistant expose settings.""" + await async_setup_component(hass, "homeassistant", {}) + registry = er.async_get(hass) + + switch_entity_entry = registry.async_get_or_create( + "switch", + "test", + "unique", + original_name="ABC", + ) + for assistant, should_expose in EXPOSE_SETTINGS.items(): + exposed_entities.async_expose_entity( + hass, assistant, switch_entity_entry.entity_id, should_expose + ) + + # Add the config entry + switch_as_x_config_entry = MockConfigEntry( + data={}, + domain=DOMAIN, + options={ + CONF_ENTITY_ID: switch_entity_entry.id, + CONF_TARGET_DOMAIN: target_domain, + }, + title="ABC", + ) + switch_as_x_config_entry.add_to_hass(hass) + + assert await hass.config_entries.async_setup(switch_as_x_config_entry.entry_id) + await hass.async_block_till_done() + + entity_entry = registry.async_get(f"{target_domain}.abc") + assert entity_entry + + # Check switch_as_x expose settings were copied from the switch + expose_settings = exposed_entities.async_get_entity_settings( + hass, entity_entry.entity_id + ) + for assistant in EXPOSE_SETTINGS: + assert expose_settings[assistant]["should_expose"] == EXPOSE_SETTINGS[assistant] + + # Check the switch is no longer exposed + expose_settings = exposed_entities.async_get_entity_settings( + hass, switch_entity_entry.entity_id + ) + for assistant in EXPOSE_SETTINGS: + assert expose_settings[assistant]["should_expose"] is False + + +@pytest.mark.parametrize("target_domain", PLATFORMS_TO_TEST) +async def test_import_expose_settings_2( + hass: HomeAssistant, + target_domain: Platform, +) -> None: + """Test importing assistant expose settings. + + This tests the expose settings are only copied from the source device when the + switch_as_x config entry is setup the first time. + """ + + await async_setup_component(hass, "homeassistant", {}) + registry = er.async_get(hass) + + switch_entity_entry = registry.async_get_or_create( + "switch", + "test", + "unique", + original_name="ABC", + ) + for assistant, should_expose in EXPOSE_SETTINGS.items(): + exposed_entities.async_expose_entity( + hass, assistant, switch_entity_entry.entity_id, should_expose + ) + + # Add the config entry + switch_as_x_config_entry = MockConfigEntry( + data={}, + domain=DOMAIN, + options={ + CONF_ENTITY_ID: switch_entity_entry.id, + CONF_TARGET_DOMAIN: target_domain, + }, + title="ABC", + ) + switch_as_x_config_entry.add_to_hass(hass) + + # Register the switch as x entity in the entity registry, this means + # the entity has been setup before + switch_as_x_entity_entry = registry.async_get_or_create( + target_domain, + "switch_as_x", + switch_as_x_config_entry.entry_id, + suggested_object_id="abc", + ) + for assistant, should_expose in EXPOSE_SETTINGS.items(): + exposed_entities.async_expose_entity( + hass, assistant, switch_as_x_entity_entry.entity_id, not should_expose + ) + + assert await hass.config_entries.async_setup(switch_as_x_config_entry.entry_id) + await hass.async_block_till_done() + + entity_entry = registry.async_get(f"{target_domain}.abc") + assert entity_entry + + # Check switch_as_x expose settings were not copied from the switch + expose_settings = exposed_entities.async_get_entity_settings( + hass, entity_entry.entity_id + ) + for assistant in EXPOSE_SETTINGS: + assert ( + expose_settings[assistant]["should_expose"] + is not EXPOSE_SETTINGS[assistant] + ) + + # Check the switch settings were not modified + expose_settings = exposed_entities.async_get_entity_settings( + hass, switch_entity_entry.entity_id + ) + for assistant in EXPOSE_SETTINGS: + assert expose_settings[assistant]["should_expose"] == EXPOSE_SETTINGS[assistant] + + +@pytest.mark.parametrize("target_domain", PLATFORMS_TO_TEST) +async def test_restore_expose_settings( + hass: HomeAssistant, + target_domain: Platform, +) -> None: + """Test removing a config entry restores assistant expose settings.""" + await async_setup_component(hass, "homeassistant", {}) + registry = er.async_get(hass) + + switch_entity_entry = registry.async_get_or_create( + "switch", + "test", + "unique", + original_name="ABC", + ) + + # Add the config entry + switch_as_x_config_entry = MockConfigEntry( + data={}, + domain=DOMAIN, + options={ + CONF_ENTITY_ID: switch_entity_entry.id, + CONF_TARGET_DOMAIN: target_domain, + }, + title="ABC", + ) + switch_as_x_config_entry.add_to_hass(hass) + + # Register the switch as x entity + switch_as_x_entity_entry = registry.async_get_or_create( + target_domain, + "switch_as_x", + switch_as_x_config_entry.entry_id, + config_entry=switch_as_x_config_entry, + suggested_object_id="abc", + ) + for assistant, should_expose in EXPOSE_SETTINGS.items(): + exposed_entities.async_expose_entity( + hass, assistant, switch_as_x_entity_entry.entity_id, should_expose + ) + + # Remove the config entry + assert await hass.config_entries.async_remove(switch_as_x_config_entry.entry_id) + await hass.async_block_till_done() + + # Check the switch expose settings were restored + expose_settings = exposed_entities.async_get_entity_settings( + hass, switch_entity_entry.entity_id + ) + for assistant in EXPOSE_SETTINGS: + assert expose_settings[assistant]["should_expose"] == EXPOSE_SETTINGS[assistant]