Add migration for old HomeWizard sensors (#105251)

Co-authored-by: Franck Nijhof <[email protected]>

homeassistant/components/homewizard/__init__.py

@@ -1,12 +1,61 @@
 """The Homewizard integration."""
 from homeassistant.config_entries import SOURCE_REAUTH, ConfigEntry
-from homeassistant.core import HomeAssistant
+from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import ConfigEntryNotReady
+from homeassistant.helpers import entity_registry as er
 
-from .const import DOMAIN, PLATFORMS
+from .const import DOMAIN, LOGGER, PLATFORMS
 from .coordinator import HWEnergyDeviceUpdateCoordinator as Coordinator
 
 
+async def _async_migrate_entries(
+    hass: HomeAssistant, config_entry: ConfigEntry
+) -> None:
+    """Migrate old entry.
+
+    The HWE-SKT had no total_power_*_kwh in 2023.11, in 2023.12 it does.
+    But simultaneously, the total_power_*_t1_kwh was removed for HWE-SKT.
+
+    This migration migrates the old unique_id to the new one, if possible.
+
+    Migration can be removed after 2024.6
+    """
+    entity_registry = er.async_get(hass)
+
+    @callback
+    def update_unique_id(entry: er.RegistryEntry) -> dict[str, str] | None:
+        replacements = {
+            "total_power_import_t1_kwh": "total_power_import_kwh",
+            "total_power_export_t1_kwh": "total_power_export_kwh",
+        }
+
+        for old_id, new_id in replacements.items():
+            if entry.unique_id.endswith(old_id):
+                new_unique_id = entry.unique_id.replace(old_id, new_id)
+                if existing_entity_id := entity_registry.async_get_entity_id(
+                    entry.domain, entry.platform, new_unique_id
+                ):
+                    LOGGER.debug(
+                        "Cannot migrate to unique_id '%s', already exists for '%s'",
+                        new_unique_id,
+                        existing_entity_id,
+                    )
+                    return None
+                LOGGER.debug(
+                    "Migrating entity '%s' unique_id from '%s' to '%s'",
+                    entry.entity_id,
+                    entry.unique_id,
+                    new_unique_id,
+                )
+                return {
+                    "new_unique_id": new_unique_id,
+                }
+
+        return None
+
+    await er.async_migrate_entries(hass, config_entry.entry_id, update_unique_id)
+
+
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up Homewizard from a config entry."""
     coordinator = Coordinator(hass)
@@ -21,6 +70,8 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
         raise
 
+    await _async_migrate_entries(hass, entry)
+
     hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator
 
     # Abort reauth config flow if active

homeassistant/components/homewizard/const.py

@@ -3,6 +3,7 @@
 
 from dataclasses import dataclass
 from datetime import timedelta
+import logging
 
 from homewizard_energy.models import Data, Device, State, System
 
@@ -11,6 +12,8 @@
 DOMAIN = "homewizard"
 PLATFORMS = [Platform.BUTTON, Platform.NUMBER, Platform.SENSOR, Platform.SWITCH]
 
+LOGGER = logging.getLogger(__package__)
+
 # Platform config.
 CONF_API_ENABLED = "api_enabled"
 CONF_DATA = "data"

homeassistant/components/homewizard/sensor.py

@@ -436,7 +436,6 @@ async def async_setup_entry(
 ) -> None:
     """Initialize sensors."""
     coordinator: HWEnergyDeviceUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
-
     async_add_entities(
         HomeWizardSensorEntity(coordinator, description)
         for description in SENSORS

tests/components/homewizard/test_init.py

@@ -7,7 +7,9 @@
 
 from homeassistant.components.homewizard.const import DOMAIN
 from homeassistant.config_entries import SOURCE_REAUTH, ConfigEntryState
+from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
+from homeassistant.helpers import entity_registry as er
 
 from tests.common import MockConfigEntry
 
@@ -118,3 +120,104 @@ async def test_load_handles_homewizardenergy_exception(
         ConfigEntryState.SETUP_RETRY,
         ConfigEntryState.SETUP_ERROR,
     )
+
+
+@pytest.mark.parametrize(
+    ("device_fixture", "old_unique_id", "new_unique_id"),
+    [
+        (
+            "HWE-SKT",
+            "aabbccddeeff_total_power_import_t1_kwh",
+            "aabbccddeeff_total_power_import_kwh",
+        ),
+        (
+            "HWE-SKT",
+            "aabbccddeeff_total_power_export_t1_kwh",
+            "aabbccddeeff_total_power_export_kwh",
+        ),
+    ],
+)
+@pytest.mark.usefixtures("mock_homewizardenergy")
+async def test_sensor_migration(
+    hass: HomeAssistant,
+    entity_registry: er.EntityRegistry,
+    mock_config_entry: MockConfigEntry,
+    old_unique_id: str,
+    new_unique_id: str,
+) -> None:
+    """Test total power T1 sensors are migrated."""
+    mock_config_entry.add_to_hass(hass)
+
+    entity: er.RegistryEntry = entity_registry.async_get_or_create(
+        domain=Platform.SENSOR,
+        platform=DOMAIN,
+        unique_id=old_unique_id,
+        config_entry=mock_config_entry,
+    )
+
+    assert entity.unique_id == old_unique_id
+
+    assert await hass.config_entries.async_setup(mock_config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    entity_migrated = entity_registry.async_get(entity.entity_id)
+    assert entity_migrated
+    assert entity_migrated.unique_id == new_unique_id
+    assert entity_migrated.previous_unique_id == old_unique_id
+
+
+@pytest.mark.parametrize(
+    ("device_fixture", "old_unique_id", "new_unique_id"),
+    [
+        (
+            "HWE-SKT",
+            "aabbccddeeff_total_power_import_t1_kwh",
+            "aabbccddeeff_total_power_import_kwh",
+        ),
+        (
+            "HWE-SKT",
+            "aabbccddeeff_total_power_export_t1_kwh",
+            "aabbccddeeff_total_power_export_kwh",
+        ),
+    ],
+)
+@pytest.mark.usefixtures("mock_homewizardenergy")
+async def test_sensor_migration_does_not_trigger(
+    hass: HomeAssistant,
+    entity_registry: er.EntityRegistry,
+    mock_config_entry: MockConfigEntry,
+    old_unique_id: str,
+    new_unique_id: str,
+) -> None:
+    """Test total power T1 sensors are not migrated when not possible."""
+    mock_config_entry.add_to_hass(hass)
+
+    old_entity: er.RegistryEntry = entity_registry.async_get_or_create(
+        domain=Platform.SENSOR,
+        platform=DOMAIN,
+        unique_id=old_unique_id,
+        config_entry=mock_config_entry,
+    )
+
+    new_entity: er.RegistryEntry = entity_registry.async_get_or_create(
+        domain=Platform.SENSOR,
+        platform=DOMAIN,
+        unique_id=new_unique_id,
+        config_entry=mock_config_entry,
+    )
+
+    assert old_entity.unique_id == old_unique_id
+    assert new_entity.unique_id == new_unique_id
+
+    assert await hass.config_entries.async_setup(mock_config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    entity = entity_registry.async_get(old_entity.entity_id)
+    assert entity
+    assert entity.unique_id == old_unique_id
+    assert entity.previous_unique_id is None
+
+    entity = entity_registry.async_get(new_entity.entity_id)
+    assert entity
+    assert entity.unique_id == new_unique_id
+    assert entity.previous_unique_id is None