add test for migration to fix #107073

homeassistant/components/swiss_public_transport/__init__.py

@@ -10,7 +10,7 @@
 from homeassistant import config_entries, core
 from homeassistant.const import Platform
 from homeassistant.exceptions import ConfigEntryError, ConfigEntryNotReady
-from homeassistant.helpers import device_registry as dr
+from homeassistant.helpers import device_registry as dr, entity_registry as er
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 
 from .const import CONF_DESTINATION, CONF_START, DOMAIN
@@ -79,23 +79,35 @@ async def async_migrate_entry(
         return False
 
     if config_entry.version == 1:
-        new = {**config_entry.data}
-
-        # Remove wrongly registered devices
+        # Remove wrongly registered devices and entries
+        new_unique_id = (
+            f"{config_entry.data[CONF_START]} {config_entry.data[CONF_DESTINATION]}"
+        )
+        entity_registry = er.async_get(hass)
         device_registry = dr.async_get(hass)
         device_entries = dr.async_entries_for_config_entry(
             device_registry, config_entry_id=config_entry.entry_id
         )
         for dev in device_entries:
             device_registry.async_remove_device(dev.id)
 
-        # Set a valid unique id for config entries
-        config_entry.unique_id = (
-            f"{config_entry.data[CONF_START]} {config_entry.data[CONF_DESTINATION]}"
+        entity_id = entity_registry.async_get_entity_id(
+            Platform.SENSOR, DOMAIN, "None_departure"
         )
+        if entity_id:
+            entity_registry.async_update_entity(
+                entity_id=entity_id,
+                new_unique_id=f"{new_unique_id}_departure",
+            )
+            _LOGGER.info(
+                "Faulty entity with unique_id 'None_departure' migrated to new unique_id '%s'",
+                f"{new_unique_id}_departure",
+            )
 
+        # Set a valid unique id for config entries
+        config_entry.unique_id = new_unique_id
         config_entry.version = 2
-        hass.config_entries.async_update_entry(config_entry, data=new)
+        hass.config_entries.async_update_entry(config_entry)
 
     _LOGGER.debug("Migration to version %s successful", config_entry.version)
 

tests/components/swiss_public_transport/test_init.py

@@ -0,0 +1,89 @@
+"""Test the swiss_public_transport config flow."""
+from unittest.mock import AsyncMock, patch
+
+from homeassistant.components.swiss_public_transport.const import (
+    CONF_DESTINATION,
+    CONF_START,
+    DOMAIN,
+)
+from homeassistant.const import Platform
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers import entity_registry as er
+
+from tests.common import MockConfigEntry
+
+MOCK_DATA_STEP = {
+    CONF_START: "test_start",
+    CONF_DESTINATION: "test_destination",
+}
+
+CONNECTIONS = [
+    {
+        "departure": "2024-01-06T18:03:00+0100",
+        "number": 0,
+        "platform": 0,
+        "transfers": 0,
+        "duration": "10",
+        "delay": 0,
+    },
+    {
+        "departure": "2024-01-06T18:04:00+0100",
+        "number": 1,
+        "platform": 1,
+        "transfers": 0,
+        "duration": "10",
+        "delay": 0,
+    },
+    {
+        "departure": "2024-01-06T18:05:00+0100",
+        "number": 2,
+        "platform": 2,
+        "transfers": 0,
+        "duration": "10",
+        "delay": 0,
+    },
+]
+
+
+def side_effect_async_get_data(x):
+    """Side effect on calling async get data which fills the connections attribute."""
+    x.connections = CONNECTIONS
+
+
+async def test_migration_1_to_2(
+    hass: HomeAssistant, entity_registry: er.EntityRegistry
+) -> None:
+    """Test successful setup."""
+
+    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",
+        )
+        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()
+        assert entity_registry.async_is_registered(
+            entity_registry.entities.get_entity_id(
+                (Platform.SENSOR, DOMAIN, "test_start test_destination_departure")
+            )
+        )
+
+        # Check change in config entry
+        assert config_entry_faulty.version == 2
+        assert config_entry_faulty.unique_id == "test_start test_destination"
+
+        # Check "None" is gone
+        assert not entity_registry.async_is_registered(
+            entity_registry.entities.get_entity_id(
+                (Platform.SENSOR, DOMAIN, "None_departure")
+            )
+        )