Teach switch_as_x about exposed entities (#92059)

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."""

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."""
+    """Unload a config entry.
-    # Unhide the wrapped entry if registered
+
+    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:
+    # Unhide the wrapped entity
-        registry.async_update_entity(entity_id, hidden_by=None)
+    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
+        )

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,15 +100,38 @@ 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)
+        def copy_custom_name(wrapped_switch: er.RegistryEntry) -> None:
-        if not wrapped_switch or wrapped_switch.name is 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)
 
+        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):
     """Represents a Switch as a ToggleEntity."""

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]