Clean-up Shelly legacy update entities (#80961)

homeassistant/components/shelly/update.py

@@ -19,9 +19,10 @@ from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.entity import EntityCategory
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
+from homeassistant.util import slugify
 
 from .const import CONF_SLEEP_PERIOD
-from .coordinator import ShellyBlockCoordinator, ShellyRpcCoordinator
+from .coordinator import ShellyBlockCoordinator, ShellyRpcCoordinator, get_entry_data
 from .entity import (
     RestEntityDescription,
     RpcEntityDescription,
@@ -30,7 +31,12 @@ from .entity import (
     async_setup_entry_rest,
     async_setup_entry_rpc,
 )
-from .utils import get_device_entry_gen
+from .utils import (
+    async_remove_shelly_entity,
+    get_block_device_name,
+    get_device_entry_gen,
+    get_rpc_device_name,
+)
 
 LOGGER = logging.getLogger(__name__)
 
@@ -117,10 +123,28 @@ async def async_setup_entry(
 ) -> None:
     """Set up update entities for Shelly component."""
     if get_device_entry_gen(config_entry) == 2:
+        # Remove legacy update binary sensor & buttons, remove in 2023.2.0
+        rpc_coordinator = get_entry_data(hass)[config_entry.entry_id].rpc
+        assert rpc_coordinator
+        mac = rpc_coordinator.mac
+        async_remove_shelly_entity(hass, "binary_sensor", f"{mac}-sys-fwupdate")
+        device_name = slugify(get_rpc_device_name(rpc_coordinator.device))
+        async_remove_shelly_entity(hass, "button", f"{device_name}_ota_update")
+        async_remove_shelly_entity(hass, "button", f"{device_name}_ota_update_beta")
+
         return async_setup_entry_rpc(
             hass, config_entry, async_add_entities, RPC_UPDATES, RpcUpdateEntity
         )
 
+    # Remove legacy update binary sensor & buttons, remove in 2023.2.0
+    block_coordinator = get_entry_data(hass)[config_entry.entry_id].block
+    assert block_coordinator
+    mac = block_coordinator.mac
+    async_remove_shelly_entity(hass, "binary_sensor", f"{mac}-fwupdate")
+    device_name = slugify(get_block_device_name(block_coordinator.device))
+    async_remove_shelly_entity(hass, "button", f"{device_name}_ota_update")
+    async_remove_shelly_entity(hass, "button", f"{device_name}_ota_update_beta")
+
     if not config_entry.data[CONF_SLEEP_PERIOD]:
         async_setup_entry_rest(
             hass,

tests/components/shelly/test_update.py

@@ -5,6 +5,8 @@ from unittest.mock import AsyncMock
 from aioshelly.exceptions import DeviceConnectionError, InvalidAuthError, RpcCallError
 import pytest
 
+from homeassistant.components.binary_sensor import DOMAIN as BINARY_SENSOR_DOMAIN
+from homeassistant.components.button import DOMAIN as BUTTON_DOMAIN
 from homeassistant.components.shelly.const import DOMAIN, REST_SENSORS_UPDATE_INTERVAL
 from homeassistant.components.update import (
     ATTR_IN_PROGRESS,
@@ -25,6 +27,44 @@ from . import MOCK_MAC, init_integration
 from tests.common import async_fire_time_changed
 
 
+@pytest.mark.parametrize(
+    "gen, domain, unique_id, object_id",
+    [
+        (1, BINARY_SENSOR_DOMAIN, f"{MOCK_MAC}-fwupdate", "firmware_update"),
+        (1, BUTTON_DOMAIN, "test_name_ota_update", "ota_update"),
+        (1, BUTTON_DOMAIN, "test_name_ota_update_beta", "ota_update_beta"),
+        (2, BINARY_SENSOR_DOMAIN, f"{MOCK_MAC}-sys-fwupdate", "firmware_update"),
+        (2, BUTTON_DOMAIN, "test_name_ota_update", "ota_update"),
+        (2, BUTTON_DOMAIN, "test_name_ota_update_beta", "ota_update_beta"),
+    ],
+)
+async def test_remove_legacy_entities(
+    hass: HomeAssistant,
+    gen,
+    domain,
+    unique_id,
+    object_id,
+    mock_block_device,
+    mock_rpc_device,
+):
+    """Test removes legacy update entities."""
+    entity_id = f"{domain}.test_name_{object_id}"
+    entity_registry = async_get(hass)
+    entity_registry.async_get_or_create(
+        domain,
+        DOMAIN,
+        unique_id,
+        suggested_object_id=f"test_name_{object_id}",
+        disabled_by=None,
+    )
+
+    assert entity_registry.async_get(entity_id) is not None
+
+    await init_integration(hass, gen)
+
+    assert entity_registry.async_get(entity_id) is None
+
+
 async def test_block_update(hass: HomeAssistant, mock_block_device, monkeypatch):
     """Test block device update entity."""
     entity_registry = async_get(hass)