Migrate sabnzbd sensors unique ids

homeassistant/components/sabnzbd/__init__.py

@@ -1,4 +1,6 @@
 """Support for monitoring an SABnzbd NZB client."""
+from __future__ import annotations
+
 from collections.abc import Callable
 import logging
 
@@ -19,7 +21,9 @@
 from homeassistant.core import HomeAssistant, ServiceCall, callback
 from homeassistant.exceptions import ConfigEntryNotReady, HomeAssistantError
 import homeassistant.helpers.config_validation as cv
+from homeassistant.helpers.device_registry import async_get
 from homeassistant.helpers.dispatcher import async_dispatcher_send
+from homeassistant.helpers.entity_registry import RegistryEntry, async_migrate_entries
 from homeassistant.helpers.event import async_track_time_interval
 from homeassistant.helpers.typing import ConfigType
 
@@ -123,8 +127,56 @@ def async_get_entry_id_for_service_call(hass: HomeAssistant, call: ServiceCall)
     raise ValueError(f"No api for API key: {call_data_api_key}")
 
 
+def update_device_identifiers(hass: HomeAssistant, entry: ConfigEntry):
+    """Update device identifiers to new identifiers."""
+    device_registry = async_get(hass)
+    device_entry = device_registry.async_get_device({(DOMAIN, DOMAIN)})
+    if device_entry and entry.entry_id in device_entry.config_entries:
+        new_identifiers = {(DOMAIN, entry.entry_id)}
+        _LOGGER.debug(
+            "Updating device id <%s> with new identifiers <%s>",
+            device_entry.id,
+            new_identifiers,
+        )
+        device_registry.async_update_device(
+            device_entry.id, new_identifiers=new_identifiers
+        )
+
+
+async def migrate_unique_id(hass: HomeAssistant, entry: ConfigEntry):
+    """Migrate entities to new unique ids (with entry_id)."""
+
+    @callback
+    def async_migrate_callback(entity_entry: RegistryEntry) -> dict | None:
+        """
+        Define a callback to migrate appropriate SabnzbdSensor entities to new unique IDs.
+
+        Old: description.key
+        New: {entry_id}_description.key
+        """
+        entry_id = entity_entry.config_entry_id
+        if entry_id is None:
+            return None
+        if entity_entry.unique_id.startswith(entry_id):
+            return None
+
+        new_unique_id = f"{entry_id}_{entity_entry.unique_id}"
+
+        _LOGGER.debug(
+            "Migrating entity %s from old unique ID '%s' to new unique ID '%s'",
+            entity_entry.entity_id,
+            entity_entry.unique_id,
+            new_unique_id,
+        )
+
+        return {"new_unique_id": new_unique_id}
+
+    await async_migrate_entries(hass, entry.entry_id, async_migrate_callback)
+
+
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up the SabNzbd Component."""
+
     sab_api = await get_client(hass, entry.data)
     if not sab_api:
         raise ConfigEntryNotReady
@@ -137,6 +189,9 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         KEY_NAME: entry.data[CONF_NAME],
     }
 
+    await migrate_unique_id(hass, entry)
+    update_device_identifiers(hass, entry)
+
     @callback
     def extract_api(func: Callable) -> Callable:
         """Define a decorator to get the correct api for a service call."""
@@ -188,6 +243,7 @@ async def async_update_sabnzbd(now):
             _LOGGER.error(err)
 
     async_track_time_interval(hass, async_update_sabnzbd, UPDATE_INTERVAL)
+
     hass.config_entries.async_setup_platforms(entry, PLATFORMS)
 
     return True

homeassistant/components/sabnzbd/sensor.py

@@ -113,11 +113,16 @@ async def async_setup_entry(
 ) -> None:
     """Set up a Sabnzbd sensor entry."""
 
-    sab_api_data = hass.data[DOMAIN][config_entry.entry_id][KEY_API_DATA]
-    client_name = hass.data[DOMAIN][config_entry.entry_id][KEY_NAME]
+    entry_id = config_entry.entry_id
+
+    sab_api_data = hass.data[DOMAIN][entry_id][KEY_API_DATA]
+    client_name = hass.data[DOMAIN][entry_id][KEY_NAME]
 
     async_add_entities(
-        [SabnzbdSensor(sab_api_data, client_name, sensor) for sensor in SENSOR_TYPES]
+        [
+            SabnzbdSensor(sab_api_data, client_name, sensor, entry_id)
+            for sensor in SENSOR_TYPES
+        ]
     )
 
 
@@ -128,17 +133,21 @@ class SabnzbdSensor(SensorEntity):
     _attr_should_poll = False
 
     def __init__(
-        self, sabnzbd_api_data, client_name, description: SabnzbdSensorEntityDescription
+        self,
+        sabnzbd_api_data,
+        client_name,
+        description: SabnzbdSensorEntityDescription,
+        entry_id,
     ):
         """Initialize the sensor."""
-        unique_id = description.key
-        self._attr_unique_id = unique_id
+
+        self._attr_unique_id = f"{entry_id}_{description.key}"
         self.entity_description = description
         self._sabnzbd_api = sabnzbd_api_data
         self._attr_name = f"{client_name} {description.name}"
         self._attr_device_info = DeviceInfo(
             entry_type=DeviceEntryType.SERVICE,
-            identifiers={(DOMAIN, DOMAIN)},
+            identifiers={(DOMAIN, entry_id)},
             name=DEFAULT_NAME,
         )
 
@@ -156,9 +165,11 @@ def update_state(self, args):
             self.entity_description.key
         )
 
-        if self.entity_description.key == SPEED_KEY:
-            self._attr_native_value = round(float(self._attr_native_value) / 1024, 1)
-        elif "size" in self.entity_description.key:
-            self._attr_native_value = round(float(self._attr_native_value), 2)
-
+        if self._attr_native_value is not None:
+            if self.entity_description.key == SPEED_KEY:
+                self._attr_native_value = round(
+                    float(self._attr_native_value) / 1024, 1
+                )
+            elif "size" in self.entity_description.key:
+                self._attr_native_value = round(float(self._attr_native_value), 2)
         self.schedule_update_ha_state()

tests/components/sabnzbd/test_config_flow.py

@@ -87,7 +87,6 @@ async def test_import_flow(hass) -> None:
         "homeassistant.components.sabnzbd.sab.SabnzbdApi.check_available",
         return_value=True,
     ):
-
         result = await hass.config_entries.flow.async_init(
             DOMAIN,
             context={"source": SOURCE_IMPORT},

tests/components/sabnzbd/test_init.py

@@ -0,0 +1,85 @@
+"""Tests for the SABnzbd Integration."""
+from unittest.mock import patch
+
+import pytest
+
+from homeassistant.components.sabnzbd import DEFAULT_NAME, DOMAIN, SENSOR_KEYS
+from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN
+from homeassistant.const import CONF_API_KEY, CONF_NAME, CONF_URL
+from homeassistant.helpers.device_registry import DeviceEntryType
+
+from tests.common import MockConfigEntry, mock_device_registry, mock_registry
+
+MOCK_ENTRY_ID = "mock_entry_id"
+
+MOCK_UNIQUE_ID = "someuniqueid"
+
+MOCK_DEVICE_ID = "somedeviceid"
+
+MOCK_DATA_VERSION_1 = {
+    CONF_API_KEY: "api_key",
+    CONF_URL: "http://127.0.0.1:8080",
+    CONF_NAME: "name",
+}
+
+MOCK_ENTRY_VERSION_1 = MockConfigEntry(
+    domain=DOMAIN, data=MOCK_DATA_VERSION_1, entry_id=MOCK_ENTRY_ID, version=1
+)
+
+
+@pytest.fixture
+def device_registry(hass):
+    """Return an empty, loaded, registry."""
+    return mock_device_registry(hass)
+
+
+@pytest.fixture
+def entity_registry(hass):
+    """Return an empty, loaded, registry."""
+    return mock_registry(hass)
+
+
+async def test_unique_id_migrate(hass, device_registry, entity_registry):
+    """Test that config flow entry is migrated correctly."""
+    # Start with the config entry at Version 1.
+    mock_entry = MOCK_ENTRY_VERSION_1
+    mock_entry.add_to_hass(hass)
+
+    mock_d_entry = device_registry.async_get_or_create(
+        config_entry_id=mock_entry.entry_id,
+        identifiers={(DOMAIN, DOMAIN)},
+        name=DEFAULT_NAME,
+        entry_type=DeviceEntryType.SERVICE,
+    )
+
+    entity_id_sensor_key = []
+
+    for sensor_key in SENSOR_KEYS:
+        mock_entity_id = f"{SENSOR_DOMAIN}.{DOMAIN}_{sensor_key}"
+        entity_registry.async_get_or_create(
+            SENSOR_DOMAIN,
+            DOMAIN,
+            unique_id=sensor_key,
+            config_entry=mock_entry,
+            device_id=mock_d_entry.id,
+        )
+        entity = entity_registry.async_get(mock_entity_id)
+        assert entity.entity_id == mock_entity_id
+        assert entity.unique_id == sensor_key
+        entity_id_sensor_key.append((mock_entity_id, sensor_key))
+
+    with patch(
+        "homeassistant.components.sabnzbd.sab.SabnzbdApi.check_available",
+        return_value=True,
+    ):
+        await hass.config_entries.async_setup(mock_entry.entry_id)
+
+        await hass.async_block_till_done()
+
+    for mock_entity_id, sensor_key in entity_id_sensor_key:
+        entity = entity_registry.async_get(mock_entity_id)
+        assert entity.unique_id == f"{MOCK_ENTRY_ID}_{sensor_key}"
+
+    assert device_registry.async_get(mock_d_entry.id).identifiers == {
+        (DOMAIN, MOCK_ENTRY_ID)
+    }