Add entity descriptions to Stookwijzer

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 @@
 
 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 @@
         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

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 .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 @@ def __init__(self, entry: StookwijzerConfigEntry) -> None:
     @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)

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,
   })
 # ---

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"
+    )