Improve migration to Uptime Kuma v2.0.0 (home-assistant#155055)

Author: tr4nt0r

homeassistant/components/uptime_kuma/coordinator.py

@@ -19,7 +19,7 @@
 from homeassistant.const import CONF_API_KEY, CONF_URL, CONF_VERIFY_SSL
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import ConfigEntryAuthFailed
-from homeassistant.helpers import entity_registry as er
+from homeassistant.helpers import device_registry as dr, entity_registry as er
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
@@ -89,7 +89,8 @@ def async_migrate_entities_unique_ids(
     """Migrate unique_ids in the entity registry after updating Uptime Kuma."""
 
     if (
-        coordinator.version is coordinator.api.version
+        coordinator.version is None
+        or coordinator.version.version == coordinator.api.version.version
         or int(coordinator.api.version.major) < 2
     ):
         return
@@ -116,6 +117,32 @@ def async_migrate_entities_unique_ids(
                 new_unique_id=f"{registry_entry.config_entry_id}_{monitor.monitor_id!s}_{registry_entry.translation_key}",
             )
 
+    # migrate device identifiers and update version
+    device_reg = dr.async_get(hass)
+    for monitor in metrics.values():
+        if device := device_reg.async_get_device(
+            {(DOMAIN, f"{coordinator.config_entry.entry_id}_{monitor.monitor_name!s}")}
+        ):
+            new_identifier = {
+                (DOMAIN, f"{coordinator.config_entry.entry_id}_{monitor.monitor_id!s}")
+            }
+            device_reg.async_update_device(
+                device.id,
+                new_identifiers=new_identifier,
+                sw_version=coordinator.api.version.version,
+            )
+    if device := device_reg.async_get_device(
+        {(DOMAIN, f"{coordinator.config_entry.entry_id}_update")}
+    ):
+        device_reg.async_update_device(
+            device.id,
+            sw_version=coordinator.api.version.version,
+        )
+
+    hass.async_create_task(
+        hass.config_entries.async_reload(coordinator.config_entry.entry_id)
+    )
+
 
 class UptimeKumaSoftwareUpdateCoordinator(DataUpdateCoordinator[LatestRelease]):
     """Uptime Kuma coordinator for retrieving update information."""

tests/components/uptime_kuma/test_sensor.py

@@ -9,10 +9,11 @@
 from pythonkuma import MonitorStatus, UptimeKumaMonitor, UptimeKumaVersion
 from syrupy.assertion import SnapshotAssertion
 
+from homeassistant.components.uptime_kuma.const import DOMAIN
 from homeassistant.config_entries import ConfigEntryState
 from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import entity_registry as er
+from homeassistant.helpers import device_registry as dr, entity_registry as er
 
 from tests.common import MockConfigEntry, async_fire_time_changed, snapshot_platform
 
@@ -53,6 +54,7 @@ async def test_migrate_unique_id(
     snapshot: SnapshotAssertion,
     entity_registry: er.EntityRegistry,
     freezer: FrozenDateTimeFactory,
+    device_registry: dr.DeviceRegistry,
 ) -> None:
     """Snapshot test states of sensor platform."""
     mock_pythonkuma.metrics.return_value = {
@@ -87,11 +89,18 @@ async def test_migrate_unique_id(
         )
     }
     mock_pythonkuma.version = UptimeKumaVersion(
-        version="2.0.0-beta.3", major="2", minor="0", patch="0-beta.3"
+        version="2.0.2", major="2", minor="0", patch="2"
     )
     freezer.tick(timedelta(seconds=30))
     async_fire_time_changed(hass)
     await hass.async_block_till_done()
 
     assert (entity := entity_registry.async_get("sensor.monitor_status"))
     assert entity.unique_id == "123456789_1_status"
+
+    assert (
+        device := device_registry.async_get_device(
+            identifiers={(DOMAIN, f"{entity.config_entry_id}_1")}
+        )
+    )
+    assert device.sw_version == "2.0.2"