diff --git a/homeassistant/components/sabnzbd/__init__.py b/homeassistant/components/sabnzbd/__init__.py index aca50e404a2..3fa5054a5f0 100644 --- a/homeassistant/components/sabnzbd/__init__.py +++ b/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.const import ( 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_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: _LOGGER.error(err) async_track_time_interval(hass, async_update_sabnzbd, UPDATE_INTERVAL) + hass.config_entries.async_setup_platforms(entry, PLATFORMS) return True diff --git a/homeassistant/components/sabnzbd/sensor.py b/homeassistant/components/sabnzbd/sensor.py index 539eaa4f097..ebdb9190ed9 100644 --- a/homeassistant/components/sabnzbd/sensor.py +++ b/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 @@ class SabnzbdSensor(SensorEntity): 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() diff --git a/tests/components/sabnzbd/test_config_flow.py b/tests/components/sabnzbd/test_config_flow.py index d04c5b18ab1..bc72dff2535 100644 --- a/tests/components/sabnzbd/test_config_flow.py +++ b/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}, diff --git a/tests/components/sabnzbd/test_init.py b/tests/components/sabnzbd/test_init.py new file mode 100644 index 00000000000..9bdef4119d0 --- /dev/null +++ b/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) + }