improve migration fix of unique ids

- follow up to home-assistant#107087
miaucl · Jan 12, 2024 · 18bc616 · 18bc616
1 parent 0257cd8
commit 18bc616
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 13 deletions.
diff --git a/homeassistant/components/swiss_public_transport/__init__.py b/homeassistant/components/swiss_public_transport/__init__.py
@@ -89,7 +89,9 @@ async def async_migrate_entry(
             device_registry, config_entry_id=config_entry.entry_id
         )
         for dev in device_entries:
-            device_registry.async_remove_device(dev.id)
+            device_registry.async_update_device(
+                dev.id, remove_config_entry_id=config_entry.entry_id
+            )
 
         entity_id = entity_registry.async_get_entity_id(
             Platform.SENSOR, DOMAIN, "None_departure"
@@ -105,12 +107,13 @@ async def async_migrate_entry(
             )
 
         # Set a valid unique id for config entries
-        config_entry.unique_id = new_unique_id
         config_entry.minor_version = 2
-        hass.config_entries.async_update_entry(config_entry)
+        hass.config_entries.async_update_entry(config_entry, unique_id=new_unique_id)
 
     _LOGGER.debug(
-        "Migration to minor version %s successful", config_entry.minor_version
+        "Migration to version %s.%s successful",
+        config_entry.version,
+        config_entry.minor_version,
     )
 
     return True
diff --git a/tests/components/swiss_public_transport/test_init.py b/tests/components/swiss_public_transport/test_init.py
@@ -45,25 +45,26 @@
 ]
 
 
-async def test_migration_1_to_2(
+async def test_migration_1_1_to_1_2(
     hass: HomeAssistant, entity_registry: er.EntityRegistry
 ) -> None:
     """Test successful setup."""
 
+    config_entry_faulty = MockConfigEntry(
+        domain=DOMAIN,
+        data=MOCK_DATA_STEP,
+        title="MIGRATION_TEST",
+        version=1,
+        minor_version=1,
+    )
+    config_entry_faulty.add_to_hass(hass)
+
     with patch(
         "homeassistant.components.swiss_public_transport.OpendataTransport",
         return_value=AsyncMock(),
     ) as mock:
         mock().connections = CONNECTIONS
 
-        config_entry_faulty = MockConfigEntry(
-            domain=DOMAIN,
-            data=MOCK_DATA_STEP,
-            title="MIGRATION_TEST",
-            minor_version=1,
-        )
-        config_entry_faulty.add_to_hass(hass)
-
         # Setup the config entry
         await hass.config_entries.async_setup(config_entry_faulty.entry_id)
         await hass.async_block_till_done()