diff --git a/homeassistant/components/stookwijzer/__init__.py b/homeassistant/components/stookwijzer/__init__.py index ee525c5323a..d8b9561bde9 100644 --- a/homeassistant/components/stookwijzer/__init__.py +++ b/homeassistant/components/stookwijzer/__init__.py @@ -2,11 +2,13 @@ from __future__ import annotations +from typing import Any + from stookwijzer import Stookwijzer from homeassistant.const import CONF_LATITUDE, CONF_LOCATION, CONF_LONGITUDE, Platform -from homeassistant.core import HomeAssistant -from homeassistant.helpers import issue_registry as ir +from homeassistant.core import HomeAssistant, callback +from homeassistant.helpers import entity_registry as er, issue_registry as ir from homeassistant.helpers.aiohttp_client import async_get_clientsession from .const import DOMAIN, LOGGER @@ -17,6 +19,8 @@ PLATFORMS = [Platform.SENSOR] async def async_setup_entry(hass: HomeAssistant, entry: StookwijzerConfigEntry) -> bool: """Set up Stookwijzer from a config entry.""" + await er.async_migrate_entries(hass, entry.entry_id, async_migrate_entity_entry) + coordinator = StookwijzerCoordinator(hass, entry) await coordinator.async_config_entry_first_refresh() @@ -71,3 +75,16 @@ async def async_migrate_entry( LOGGER.debug("Migration to version %s successful", entry.version) return True + + +@callback +def async_migrate_entity_entry(entity_entry: er.RegistryEntry) -> dict[str, Any] | None: + """Migrate Stookwijzer entity entries. + + - Migrates unique ID for the old Stookwijzer sensors to the new unique ID. + """ + if entity_entry.unique_id == entity_entry.config_entry_id: + return {"new_unique_id": f"{entity_entry.config_entry_id}_advice"} + + # No migration needed + return None diff --git a/homeassistant/components/stookwijzer/sensor.py b/homeassistant/components/stookwijzer/sensor.py index 25396639ecd..1027f8c0e98 100644 --- a/homeassistant/components/stookwijzer/sensor.py +++ b/homeassistant/components/stookwijzer/sensor.py @@ -2,7 +2,16 @@ from __future__ import annotations -from homeassistant.components.sensor import SensorDeviceClass, SensorEntity +from collections.abc import Callable +from dataclasses import dataclass + +from stookwijzer import Stookwijzer + +from homeassistant.components.sensor import ( + SensorDeviceClass, + SensorEntity, + SensorEntityDescription, +) from homeassistant.core import HomeAssistant from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo from homeassistant.helpers.entity_platform import AddEntitiesCallback @@ -12,29 +21,51 @@ from .const import DOMAIN from .coordinator import StookwijzerConfigEntry, StookwijzerCoordinator +@dataclass(kw_only=True, frozen=True) +class StookwijzerSensorDescription(SensorEntityDescription): + """Class describing Stookwijzer sensor entities.""" + + value_fn: Callable[[Stookwijzer], str | None] + + +STOOKWIJZER_SENSORS = [ + StookwijzerSensorDescription( + key="advice", + translation_key="advice", + device_class=SensorDeviceClass.ENUM, + value_fn=lambda client: client.advice, + options=["code_yellow", "code_orange", "code_red"], + ), +] + + async def async_setup_entry( hass: HomeAssistant, entry: StookwijzerConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up Stookwijzer sensor from a config entry.""" - async_add_entities([StookwijzerSensor(entry)]) + async_add_entities( + StookwijzerSensor(description, entry) for description in STOOKWIJZER_SENSORS + ) class StookwijzerSensor(CoordinatorEntity[StookwijzerCoordinator], SensorEntity): """Defines a Stookwijzer binary sensor.""" + entity_description: StookwijzerSensorDescription _attr_attribution = "Data provided by atlasleefomgeving.nl" - _attr_device_class = SensorDeviceClass.ENUM _attr_has_entity_name = True - _attr_translation_key = "advice" - def __init__(self, entry: StookwijzerConfigEntry) -> None: + def __init__( + self, + description: StookwijzerSensorDescription, + entry: StookwijzerConfigEntry, + ) -> None: """Initialize a Stookwijzer device.""" super().__init__(entry.runtime_data) - self._client = entry.runtime_data.client - self._attr_options = ["code_yellow", "code_orange", "code_red"] - self._attr_unique_id = entry.entry_id + self.entity_description = description + self._attr_unique_id = f"{entry.entry_id}_{description.key}" self._attr_device_info = DeviceInfo( identifiers={(DOMAIN, entry.entry_id)}, manufacturer="Atlas Leefomgeving", @@ -45,4 +76,4 @@ class StookwijzerSensor(CoordinatorEntity[StookwijzerCoordinator], SensorEntity) @property def native_value(self) -> str | None: """Return the state of the device.""" - return self._client.advice + return self.entity_description.value_fn(self.coordinator.client) diff --git a/tests/components/stookwijzer/snapshots/test_sensor.ambr b/tests/components/stookwijzer/snapshots/test_sensor.ambr index bd7fcc6fef5..195f3498225 100644 --- a/tests/components/stookwijzer/snapshots/test_sensor.ambr +++ b/tests/components/stookwijzer/snapshots/test_sensor.ambr @@ -34,7 +34,7 @@ 'previous_unique_id': None, 'supported_features': 0, 'translation_key': 'advice', - 'unique_id': '12345', + 'unique_id': '12345_advice', 'unit_of_measurement': None, }) # --- diff --git a/tests/components/stookwijzer/test_init.py b/tests/components/stookwijzer/test_init.py index 1a84954c21a..0774def7813 100644 --- a/tests/components/stookwijzer/test_init.py +++ b/tests/components/stookwijzer/test_init.py @@ -2,11 +2,14 @@ from unittest.mock import MagicMock +import pytest + +from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN from homeassistant.components.stookwijzer.const import DOMAIN from homeassistant.config_entries import ConfigEntryState from homeassistant.const import CONF_LATITUDE, CONF_LONGITUDE 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 tests.common import MockConfigEntry @@ -53,3 +56,44 @@ async def test_entry_migration_failure( assert issue_registry.async_get_issue(DOMAIN, "location_migration_failed") assert len(mock_stookwijzer.async_transform_coordinates.mock_calls) == 1 + + +@pytest.mark.usefixtures("mock_stookwijzer") +async def test_entity_entry_migration( + hass: HomeAssistant, + mock_config_entry: MockConfigEntry, + entity_registry: er.EntityRegistry, +) -> None: + """Test successful migration of entry data.""" + entity = entity_registry.async_get_or_create( + suggested_object_id="advice", + disabled_by=None, + domain=SENSOR_DOMAIN, + platform=DOMAIN, + unique_id=mock_config_entry.entry_id, + config_entry=mock_config_entry, + ) + + assert entity.unique_id == mock_config_entry.entry_id + + mock_config_entry.add_to_hass(hass) + await hass.config_entries.async_setup(mock_config_entry.entry_id) + await hass.async_block_till_done() + + assert ( + entity_registry.async_get_entity_id( + SENSOR_DOMAIN, + DOMAIN, + mock_config_entry.entry_id, + ) + is None + ) + + assert ( + entity_registry.async_get_entity_id( + SENSOR_DOMAIN, + DOMAIN, + f"{mock_config_entry.entry_id}_advice", + ) + == "sensor.advice" + )