From d84d83a42ae1718cdc21212503a97483fe048df0 Mon Sep 17 00:00:00 2001 From: Joost Lekkerkerker Date: Fri, 29 Sep 2023 19:30:35 +0200 Subject: [PATCH] Migrate WAQI unique id (#101112) * Migrate unique_id * Add docstring --- homeassistant/components/waqi/__init__.py | 16 +++++++++++++ homeassistant/components/waqi/sensor.py | 2 +- tests/components/waqi/test_sensor.py | 29 ++++++++++++++++++++++- 3 files changed, 45 insertions(+), 2 deletions(-) diff --git a/homeassistant/components/waqi/__init__.py b/homeassistant/components/waqi/__init__.py index bc51a91364c..d3cf1af21a2 100644 --- a/homeassistant/components/waqi/__init__.py +++ b/homeassistant/components/waqi/__init__.py @@ -7,6 +7,7 @@ from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_API_KEY, Platform from homeassistant.core import HomeAssistant from homeassistant.helpers.aiohttp_client import async_get_clientsession +import homeassistant.helpers.entity_registry as er from .const import DOMAIN from .coordinator import WAQIDataUpdateCoordinator @@ -17,6 +18,8 @@ PLATFORMS: list[Platform] = [Platform.SENSOR] async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Set up World Air Quality Index (WAQI) from a config entry.""" + await _migrate_unique_ids(hass, entry) + client = WAQIClient(session=async_get_clientsession(hass)) client.authenticate(entry.data[CONF_API_KEY]) @@ -35,3 +38,16 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: hass.data[DOMAIN].pop(entry.entry_id) return unload_ok + + +async def _migrate_unique_ids(hass: HomeAssistant, entry: ConfigEntry) -> None: + """Migrate pre-config flow unique ids.""" + entity_registry = er.async_get(hass) + registry_entries = er.async_entries_for_config_entry( + entity_registry, entry.entry_id + ) + for reg_entry in registry_entries: + if isinstance(reg_entry.unique_id, int): + entity_registry.async_update_entity( + reg_entry.entity_id, new_unique_id=f"{reg_entry.unique_id}_air_quality" + ) diff --git a/homeassistant/components/waqi/sensor.py b/homeassistant/components/waqi/sensor.py index 0ad295ca5af..62170b329f4 100644 --- a/homeassistant/components/waqi/sensor.py +++ b/homeassistant/components/waqi/sensor.py @@ -159,7 +159,7 @@ class WaqiSensor(CoordinatorEntity[WAQIDataUpdateCoordinator], SensorEntity): """Initialize the sensor.""" super().__init__(coordinator) self._attr_name = f"WAQI {self.coordinator.data.city.name}" - self._attr_unique_id = str(coordinator.data.station_id) + self._attr_unique_id = f"{coordinator.data.station_id}_air_quality" @property def native_value(self) -> int | None: diff --git a/tests/components/waqi/test_sensor.py b/tests/components/waqi/test_sensor.py index ef434bcc544..7feb37a1b09 100644 --- a/tests/components/waqi/test_sensor.py +++ b/tests/components/waqi/test_sensor.py @@ -4,6 +4,7 @@ from unittest.mock import patch from aiowaqi import WAQIAirQuality, WAQIError, WAQISearchResult +from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN from homeassistant.components.waqi.const import CONF_STATION_NUMBER, DOMAIN from homeassistant.components.waqi.sensor import CONF_LOCATIONS, CONF_STATIONS from homeassistant.config_entries import SOURCE_IMPORT, ConfigEntryState @@ -15,7 +16,7 @@ from homeassistant.const import ( Platform, ) from homeassistant.core import HomeAssistant -from homeassistant.helpers import issue_registry as ir +from homeassistant.helpers import entity_registry as er, issue_registry as ir from homeassistant.setup import async_setup_component from tests.common import MockConfigEntry, load_fixture @@ -93,6 +94,32 @@ async def test_legacy_migration_already_imported( assert len(issue_registry.issues) == 1 +async def test_sensor_id_migration( + hass: HomeAssistant, mock_config_entry: MockConfigEntry +) -> None: + """Test migrating unique id for original sensor.""" + mock_config_entry.add_to_hass(hass) + entity_registry = er.async_get(hass) + entity_registry.async_get_or_create( + SENSOR_DOMAIN, DOMAIN, 4584, config_entry=mock_config_entry + ) + with patch( + "aiowaqi.WAQIClient.get_by_station_number", + return_value=WAQIAirQuality.from_dict( + json.loads(load_fixture("waqi/air_quality_sensor.json")) + ), + ): + assert await async_setup_component(hass, DOMAIN, {}) + await hass.async_block_till_done() + entities = er.async_entries_for_config_entry( + entity_registry, mock_config_entry.entry_id + ) + assert len(entities) == 1 + assert hass.states.get("sensor.waqi_4584") + assert hass.states.get("sensor.waqi_de_jongweg_utrecht") is None + assert entities[0].unique_id == "4584_air_quality" + + async def test_sensor(hass: HomeAssistant, mock_config_entry: MockConfigEntry) -> None: """Test failed update.""" mock_config_entry.add_to_hass(hass)