Restore the Home Connect program option entities (#156401)

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

homeassistant/components/home_connect/binary_sensor.py

@@ -169,6 +169,7 @@ async def async_setup_entry(
 ) -> None:
     """Set up the Home Connect binary sensor."""
     setup_home_connect_entry(
+        hass,
         entry,
         _get_entities_for_appliance,
         async_add_entities,

homeassistant/components/home_connect/button.py

@@ -73,6 +73,7 @@ async def async_setup_entry(
 ) -> None:
     """Set up the Home Connect button entities."""
     setup_home_connect_entry(
+        hass,
         entry,
         _get_entities_for_appliance,
         async_add_entities,

homeassistant/components/home_connect/common.py

@@ -7,18 +7,44 @@ from typing import cast
 
 from aiohomeconnect.model import EventKey
 
+from homeassistant.const import Platform
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers import entity_registry as er
+from homeassistant.helpers.entity import EntityDescription
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
+from .const import DOMAIN
 from .coordinator import HomeConnectApplianceData, HomeConnectConfigEntry
 from .entity import HomeConnectEntity, HomeConnectOptionEntity
 
 
+def should_add_option_entity(
+    description: EntityDescription,
+    appliance: HomeConnectApplianceData,
+    entity_registry: er.EntityRegistry,
+    platform: Platform,
+) -> bool:
+    """Check if the option entity should be added for the appliance.
+
+    This function returns `True` if the option is available in the appliance options
+    or if the entity was added in previous loads of this integration.
+    """
+    description_key = description.key
+    return description_key in appliance.options or (
+        entity_registry.async_get_entity_id(
+            platform, DOMAIN, f"{appliance.info.ha_id}-{description_key}"
+        )
+        is not None
+    )
+
+
 def _create_option_entities(
+    entity_registry: er.EntityRegistry,
     entry: HomeConnectConfigEntry,
     appliance: HomeConnectApplianceData,
     known_entity_unique_ids: dict[str, str],
     get_option_entities_for_appliance: Callable[
-        [HomeConnectConfigEntry, HomeConnectApplianceData],
+        [HomeConnectConfigEntry, HomeConnectApplianceData, er.EntityRegistry],
         list[HomeConnectOptionEntity],
     ],
     async_add_entities: AddConfigEntryEntitiesCallback,
@@ -26,7 +52,9 @@ def _create_option_entities(
     """Create the required option entities for the appliances."""
     option_entities_to_add = [
         entity
-        for entity in get_option_entities_for_appliance(entry, appliance)
+        for entity in get_option_entities_for_appliance(
+            entry, appliance, entity_registry
+        )
         if entity.unique_id not in known_entity_unique_ids
     ]
     known_entity_unique_ids.update(
@@ -39,13 +67,14 @@ def _create_option_entities(
 
 
 def _handle_paired_or_connected_appliance(
+    hass: HomeAssistant,
     entry: HomeConnectConfigEntry,
     known_entity_unique_ids: dict[str, str],
     get_entities_for_appliance: Callable[
         [HomeConnectConfigEntry, HomeConnectApplianceData], list[HomeConnectEntity]
     ],
     get_option_entities_for_appliance: Callable[
-        [HomeConnectConfigEntry, HomeConnectApplianceData],
+        [HomeConnectConfigEntry, HomeConnectApplianceData, er.EntityRegistry],
         list[HomeConnectOptionEntity],
     ]
     | None,
@@ -60,6 +89,7 @@ def _handle_paired_or_connected_appliance(
     already or it is the first time we see them when the appliance is connected.
     """
     entities: list[HomeConnectEntity] = []
+    entity_registry = er.async_get(hass)
     for appliance in entry.runtime_data.data.values():
         entities_to_add = [
             entity
@@ -69,7 +99,9 @@ def _handle_paired_or_connected_appliance(
         if get_option_entities_for_appliance:
             entities_to_add.extend(
                 entity
-                for entity in get_option_entities_for_appliance(entry, appliance)
+                for entity in get_option_entities_for_appliance(
+                    entry, appliance, entity_registry
+                )
                 if entity.unique_id not in known_entity_unique_ids
             )
             for event_key in (
@@ -80,6 +112,7 @@ def _handle_paired_or_connected_appliance(
                     entry.runtime_data.async_add_listener(
                         partial(
                             _create_option_entities,
+                            entity_registry,
                             entry,
                             appliance,
                             known_entity_unique_ids,
@@ -120,13 +153,14 @@ def _handle_depaired_appliance(
 
 
 def setup_home_connect_entry(
+    hass: HomeAssistant,
     entry: HomeConnectConfigEntry,
     get_entities_for_appliance: Callable[
         [HomeConnectConfigEntry, HomeConnectApplianceData], list[HomeConnectEntity]
     ],
     async_add_entities: AddConfigEntryEntitiesCallback,
     get_option_entities_for_appliance: Callable[
-        [HomeConnectConfigEntry, HomeConnectApplianceData],
+        [HomeConnectConfigEntry, HomeConnectApplianceData, er.EntityRegistry],
         list[HomeConnectOptionEntity],
     ]
     | None = None,
@@ -141,6 +175,7 @@ def setup_home_connect_entry(
         entry.runtime_data.async_add_special_listener(
             partial(
                 _handle_paired_or_connected_appliance,
+                hass,
                 entry,
                 known_entity_unique_ids,
                 get_entities_for_appliance,

homeassistant/components/home_connect/light.py

@@ -96,6 +96,7 @@ async def async_setup_entry(
 ) -> None:
     """Set up the Home Connect light."""
     setup_home_connect_entry(
+        hass,
         entry,
         _get_entities_for_appliance,
         async_add_entities,

homeassistant/components/home_connect/number.py

@@ -11,12 +11,13 @@ from homeassistant.components.number import (
     NumberEntity,
     NumberEntityDescription,
 )
-from homeassistant.const import PERCENTAGE
+from homeassistant.const import PERCENTAGE, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
+from homeassistant.helpers import entity_registry as er
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
-from .common import setup_home_connect_entry
+from .common import setup_home_connect_entry, should_add_option_entity
 from .const import DOMAIN, UNIT_MAP
 from .coordinator import HomeConnectApplianceData, HomeConnectConfigEntry
 from .entity import HomeConnectEntity, HomeConnectOptionEntity, constraint_fetcher
@@ -136,12 +137,15 @@ def _get_entities_for_appliance(
 def _get_option_entities_for_appliance(
     entry: HomeConnectConfigEntry,
     appliance: HomeConnectApplianceData,
+    entity_registry: er.EntityRegistry,
 ) -> list[HomeConnectOptionEntity]:
     """Get a list of currently available option entities."""
     return [
         HomeConnectOptionNumberEntity(entry.runtime_data, appliance, description)
         for description in NUMBER_OPTIONS
-        if description.key in appliance.options
+        if should_add_option_entity(
+            description, appliance, entity_registry, Platform.NUMBER
+        )
     ]
 
 
@@ -152,6 +156,7 @@ async def async_setup_entry(
 ) -> None:
     """Set up the Home Connect number."""
     setup_home_connect_entry(
+        hass,
         entry,
         _get_entities_for_appliance,
         async_add_entities,

homeassistant/components/home_connect/select.py

@@ -11,11 +11,13 @@ from aiohomeconnect.model.error import HomeConnectError
 from aiohomeconnect.model.program import Execution
 
 from homeassistant.components.select import SelectEntity, SelectEntityDescription
+from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import HomeAssistantError
+from homeassistant.helpers import entity_registry as er
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
-from .common import setup_home_connect_entry
+from .common import setup_home_connect_entry, should_add_option_entity
 from .const import (
     AVAILABLE_MAPS_ENUM,
     BEAN_AMOUNT_OPTIONS,
@@ -358,12 +360,13 @@ def _get_entities_for_appliance(
 def _get_option_entities_for_appliance(
     entry: HomeConnectConfigEntry,
     appliance: HomeConnectApplianceData,
+    entity_registry: er.EntityRegistry,
 ) -> list[HomeConnectOptionEntity]:
     """Get a list of entities."""
     return [
         HomeConnectSelectOptionEntity(entry.runtime_data, appliance, desc)
         for desc in PROGRAM_SELECT_OPTION_ENTITY_DESCRIPTIONS
-        if desc.key in appliance.options
+        if should_add_option_entity(desc, appliance, entity_registry, Platform.SELECT)
     ]
 
 
@@ -374,6 +377,7 @@ async def async_setup_entry(
 ) -> None:
     """Set up the Home Connect select entities."""
     setup_home_connect_entry(
+        hass,
         entry,
         _get_entities_for_appliance,
         async_add_entities,

homeassistant/components/home_connect/sensor.py

@@ -540,6 +540,7 @@ async def async_setup_entry(
 ) -> None:
     """Set up the Home Connect sensor."""
     setup_home_connect_entry(
+        hass,
         entry,
         _get_entities_for_appliance,
         async_add_entities,

homeassistant/components/home_connect/switch.py

@@ -7,12 +7,14 @@ from aiohomeconnect.model import OptionKey, SettingKey
 from aiohomeconnect.model.error import HomeConnectError
 
 from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription
+from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
+from homeassistant.helpers import entity_registry as er
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 from homeassistant.helpers.typing import UNDEFINED, UndefinedType
 
-from .common import setup_home_connect_entry
+from .common import setup_home_connect_entry, should_add_option_entity
 from .const import BSH_POWER_OFF, BSH_POWER_ON, BSH_POWER_STANDBY, DOMAIN
 from .coordinator import HomeConnectApplianceData, HomeConnectConfigEntry
 from .entity import HomeConnectEntity, HomeConnectOptionEntity
@@ -190,12 +192,15 @@ def _get_entities_for_appliance(
 def _get_option_entities_for_appliance(
     entry: HomeConnectConfigEntry,
     appliance: HomeConnectApplianceData,
+    entity_registry: er.EntityRegistry,
 ) -> list[HomeConnectOptionEntity]:
     """Get a list of currently available option entities."""
     return [
         HomeConnectSwitchOptionEntity(entry.runtime_data, appliance, description)
         for description in SWITCH_OPTIONS
-        if description.key in appliance.options
+        if should_add_option_entity(
+            description, appliance, entity_registry, Platform.SWITCH
+        )
     ]
 
 
@@ -206,6 +211,7 @@ async def async_setup_entry(
 ) -> None:
     """Set up the Home Connect switch."""
     setup_home_connect_entry(
+        hass,
         entry,
         _get_entities_for_appliance,
         async_add_entities,

tests/components/home_connect/test_number.py

@@ -44,7 +44,12 @@ from homeassistant.components.number import (
     SERVICE_SET_VALUE,
 )
 from homeassistant.config_entries import ConfigEntryState
-from homeassistant.const import ATTR_ENTITY_ID, STATE_UNAVAILABLE, Platform
+from homeassistant.const import (
+    ATTR_ENTITY_ID,
+    ATTR_RESTORED,
+    STATE_UNAVAILABLE,
+    Platform,
+)
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers import device_registry as dr, entity_registry as er
@@ -757,3 +762,42 @@ async def test_options_available_when_program_is_null(
     state = hass.states.get(entity_id)
     assert state
     assert state.state != STATE_UNAVAILABLE
+
+
+@pytest.mark.parametrize("appliance", ["Oven"], indirect=True)
+async def test_restore_option_entity(
+    hass: HomeAssistant,
+    entity_registry: er.EntityRegistry,
+    client: MagicMock,
+    config_entry: MockConfigEntry,
+    integration_setup: Callable[[MagicMock], Awaitable[bool]],
+    appliance: HomeAppliance,
+) -> None:
+    """Test restoration of option entities when program options are missing.
+
+    This test ensures that number entities representing options are restored
+    to the entity registry and set to unavailable if the current available
+    program does not include them, but they existed previously.
+    """
+    entity_id = "number.oven_setpoint_temperature"
+    client.get_available_program = AsyncMock(
+        return_value=ProgramDefinition(
+            ProgramKey.UNKNOWN,
+            options=[],
+        )
+    )
+
+    entity_registry.async_get_or_create(
+        Platform.NUMBER,
+        DOMAIN,
+        f"{appliance.ha_id}-{OptionKey.COOKING_OVEN_SETPOINT_TEMPERATURE}",
+        suggested_object_id="oven_setpoint_temperature",
+    )
+
+    assert await integration_setup(client)
+    assert config_entry.state is ConfigEntryState.LOADED
+
+    state = hass.states.get(entity_id)
+    assert state is not None
+    assert state.state == STATE_UNAVAILABLE
+    assert not state.attributes.get(ATTR_RESTORED)

tests/components/home_connect/test_select.py

@@ -44,6 +44,7 @@ from homeassistant.components.select import (
 from homeassistant.config_entries import ConfigEntryState
 from homeassistant.const import (
     ATTR_ENTITY_ID,
+    ATTR_RESTORED,
     SERVICE_SELECT_OPTION,
     STATE_UNAVAILABLE,
     STATE_UNKNOWN,
@@ -1105,3 +1106,42 @@ async def test_options_available_when_program_is_null(
     state = hass.states.get(entity_id)
     assert state
     assert state.state != STATE_UNAVAILABLE
+
+
+@pytest.mark.parametrize("appliance", ["Washer"], indirect=True)
+async def test_restore_option_entity(
+    hass: HomeAssistant,
+    entity_registry: er.EntityRegistry,
+    client: MagicMock,
+    config_entry: MockConfigEntry,
+    integration_setup: Callable[[MagicMock], Awaitable[bool]],
+    appliance: HomeAppliance,
+) -> None:
+    """Test restoration of option entities when program options are missing.
+
+    This test ensures that number entities representing options are restored
+    to the entity registry and set to unavailable if the current available
+    program does not include them, but they existed previously.
+    """
+    entity_id = "select.washer_temperature"
+    client.get_available_program = AsyncMock(
+        return_value=ProgramDefinition(
+            ProgramKey.UNKNOWN,
+            options=[],
+        )
+    )
+
+    entity_registry.async_get_or_create(
+        Platform.SELECT,
+        DOMAIN,
+        f"{appliance.ha_id}-{OptionKey.LAUNDRY_CARE_WASHER_TEMPERATURE}",
+        suggested_object_id="washer_temperature",
+    )
+
+    assert await integration_setup(client)
+    assert config_entry.state is ConfigEntryState.LOADED
+
+    state = hass.states.get(entity_id)
+    assert state is not None
+    assert state.state == STATE_UNAVAILABLE
+    assert not state.attributes.get(ATTR_RESTORED)

tests/components/home_connect/test_switch.py

@@ -38,6 +38,7 @@ from homeassistant.components.switch import DOMAIN as SWITCH_DOMAIN
 from homeassistant.config_entries import ConfigEntryState
 from homeassistant.const import (
     ATTR_ENTITY_ID,
+    ATTR_RESTORED,
     SERVICE_TURN_OFF,
     SERVICE_TURN_ON,
     STATE_OFF,
@@ -878,3 +879,42 @@ async def test_options_available_when_program_is_null(
     state = hass.states.get(entity_id)
     assert state
     assert state.state != STATE_UNAVAILABLE
+
+
+@pytest.mark.parametrize("appliance", ["Dishwasher"], indirect=True)
+async def test_restore_option_entity(
+    hass: HomeAssistant,
+    entity_registry: er.EntityRegistry,
+    client: MagicMock,
+    config_entry: MockConfigEntry,
+    integration_setup: Callable[[MagicMock], Awaitable[bool]],
+    appliance: HomeAppliance,
+) -> None:
+    """Test restoration of option entities when program options are missing.
+
+    This test ensures that number entities representing options are restored
+    to the entity registry and set to unavailable if the current available
+    program does not include them, but they existed previously.
+    """
+    entity_id = "switch.dishwasher_half_load"
+    client.get_available_program = AsyncMock(
+        return_value=ProgramDefinition(
+            ProgramKey.UNKNOWN,
+            options=[],
+        )
+    )
+
+    entity_registry.async_get_or_create(
+        Platform.SWITCH,
+        DOMAIN,
+        f"{appliance.ha_id}-{OptionKey.DISHCARE_DISHWASHER_HALF_LOAD}",
+        suggested_object_id="dishwasher_half_load",
+    )
+
+    assert await integration_setup(client)
+    assert config_entry.state is ConfigEntryState.LOADED
+
+    state = hass.states.get(entity_id)
+    assert state is not None
+    assert state.state == STATE_UNAVAILABLE
+    assert not state.attributes.get(ATTR_RESTORED)