Add Shelly removal condition for virtual components (home-assistant#152312)

Author: thecode

homeassistant/components/shelly/binary_sensor.py

@@ -40,6 +40,7 @@
     get_virtual_component_ids,
     is_block_momentary_input,
     is_rpc_momentary_input,
+    is_view_for_platform,
 )
 
 PARALLEL_UPDATES = 0
@@ -273,6 +274,9 @@ def __init__(
     "boolean": RpcBinarySensorDescription(
         key="boolean",
         sub_key="value",
+        removal_condition=lambda config, _status, key: not is_view_for_platform(
+            config, key, BINARY_SENSOR_PLATFORM
+        ),
     ),
     "calibration": RpcBinarySensorDescription(
         key="blutrv",

homeassistant/components/shelly/number.py

@@ -41,6 +41,7 @@
     get_device_entry_gen,
     get_virtual_component_ids,
     get_virtual_component_unit,
+    is_view_for_platform,
 )
 
 PARALLEL_UPDATES = 0
@@ -185,6 +186,9 @@ async def async_set_native_value(self, value: float) -> None:
     "number": RpcNumberDescription(
         key="number",
         sub_key="value",
+        removal_condition=lambda config, _status, key: not is_view_for_platform(
+            config, key, NUMBER_PLATFORM
+        ),
         max_fn=lambda config: config["max"],
         min_fn=lambda config: config["min"],
         mode_fn=lambda config: VIRTUAL_NUMBER_MODE_MAP.get(

homeassistant/components/shelly/select.py

@@ -26,6 +26,7 @@
     async_remove_orphaned_entities,
     get_device_entry_gen,
     get_virtual_component_ids,
+    is_view_for_platform,
 )
 
 PARALLEL_UPDATES = 0
@@ -40,6 +41,9 @@ class RpcSelectDescription(RpcEntityDescription, SelectEntityDescription):
     "enum": RpcSelectDescription(
         key="enum",
         sub_key="value",
+        removal_condition=lambda config, _status, key: not is_view_for_platform(
+            config, key, SELECT_PLATFORM
+        ),
     ),
 }
 

homeassistant/components/shelly/sensor.py

@@ -63,6 +63,7 @@
     get_virtual_component_ids,
     get_virtual_component_unit,
     is_rpc_wifi_stations_disabled,
+    is_view_for_platform,
 )
 
 PARALLEL_UPDATES = 0
@@ -1385,10 +1386,16 @@ def __init__(
     "text": RpcSensorDescription(
         key="text",
         sub_key="value",
+        removal_condition=lambda config, _status, key: not is_view_for_platform(
+            config, key, SENSOR_PLATFORM
+        ),
     ),
     "number": RpcSensorDescription(
         key="number",
         sub_key="value",
+        removal_condition=lambda config, _status, key: not is_view_for_platform(
+            config, key, SENSOR_PLATFORM
+        ),
         unit=get_virtual_component_unit,
         device_class_fn=lambda config: ROLE_TO_DEVICE_CLASS_MAP.get(config["role"])
         if "role" in config
@@ -1397,6 +1404,9 @@ def __init__(
     "enum": RpcSensorDescription(
         key="enum",
         sub_key="value",
+        removal_condition=lambda config, _status, key: not is_view_for_platform(
+            config, key, SENSOR_PLATFORM
+        ),
         options_fn=lambda config: config["options"],
         device_class=SensorDeviceClass.ENUM,
     ),

homeassistant/components/shelly/switch.py

@@ -37,6 +37,7 @@
     get_virtual_component_ids,
     is_block_exclude_from_relay,
     is_rpc_exclude_from_relay,
+    is_view_for_platform,
 )
 
 PARALLEL_UPDATES = 0
@@ -89,6 +90,9 @@ class RpcSwitchDescription(RpcEntityDescription, SwitchEntityDescription):
     "boolean": RpcSwitchDescription(
         key="boolean",
         sub_key="value",
+        removal_condition=lambda config, _status, key: not is_view_for_platform(
+            config, key, SWITCH_PLATFORM
+        ),
         is_on=lambda status: bool(status["value"]),
         method_on="Boolean.Set",
         method_off="Boolean.Set",

homeassistant/components/shelly/text.py

@@ -26,6 +26,7 @@
     async_remove_orphaned_entities,
     get_device_entry_gen,
     get_virtual_component_ids,
+    is_view_for_platform,
 )
 
 PARALLEL_UPDATES = 0
@@ -40,6 +41,9 @@ class RpcTextDescription(RpcEntityDescription, TextEntityDescription):
     "text": RpcTextDescription(
         key="text",
         sub_key="value",
+        removal_condition=lambda config, _status, key: not is_view_for_platform(
+            config, key, TEXT_PLATFORM
+        ),
     ),
 }
 

homeassistant/components/shelly/utils.py

@@ -654,6 +654,13 @@ def get_virtual_component_ids(config: dict[str, Any], platform: str) -> list[str
     return ids
 
 
+def is_view_for_platform(config: dict[str, Any], key: str, platform: str) -> bool:
+    """Return true if the virtual component view match the platform."""
+    component = VIRTUAL_COMPONENTS_MAP[platform]
+    view = config[key]["meta"]["ui"]["view"]
+    return view in component["modes"]
+
+
 def get_virtual_component_unit(config: dict[str, Any]) -> str | None:
     """Return the unit of a virtual component.
 

tests/components/shelly/conftest.py

@@ -762,3 +762,16 @@ def mock_setup() -> Generator[AsyncMock]:
         "homeassistant.components.shelly.async_setup", return_value=True
     ) as mock_setup:
         yield mock_setup
+
+
+@pytest.fixture
+def disable_async_remove_shelly_rpc_entities() -> Generator[None]:
+    """Patch out async_remove_shelly_rpc_entities.
+
+    This is used by virtual components tests that should not create entities,
+    without it async_remove_shelly_rpc_entities will clean up the entities.
+    """
+    with patch(
+        "homeassistant.components.shelly.utils.async_remove_shelly_rpc_entities"
+    ):
+        yield

tests/components/shelly/test_binary_sensor.py

@@ -435,6 +435,7 @@ async def test_rpc_device_virtual_binary_sensor(
     assert state.state == STATE_OFF
 
 
+@pytest.mark.usefixtures("disable_async_remove_shelly_rpc_entities")
 async def test_rpc_remove_virtual_binary_sensor_when_mode_toggle(
     hass: HomeAssistant,
     entity_registry: EntityRegistry,

tests/components/shelly/test_number.py

@@ -340,6 +340,7 @@ async def test_rpc_device_virtual_number(
     assert state.state == "56.7"
 
 
+@pytest.mark.usefixtures("disable_async_remove_shelly_rpc_entities")
 async def test_rpc_remove_virtual_number_when_mode_label(
     hass: HomeAssistant,
     entity_registry: EntityRegistry,