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
@@ -33,6 +33,7 @@ async def async_setup_entry(
 
     session = async_get_clientsession(hass)
     opendata = OpendataTransport(start, destination, session)
+    _LOGGER.info(opendata)
 
     try:
         await opendata.async_get_data()
@@ -79,23 +80,30 @@ 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_remove(entity_id)
 
+        # 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)
+
+        # await hass.config_entries.async_forward_entry_setups(config_entry, PLATFORMS)
 
     _LOGGER.debug("Migration to version %s successful", config_entry.version)
 

tests/components/swiss_public_transport/test_migration.py

@@ -0,0 +1,111 @@
+"""Test the swiss_public_transport config flow."""
+from unittest.mock import patch
+
+from homeassistant.components.swiss_public_transport import (
+    async_migrate_entry,
+    async_setup_entry,
+)
+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",
+}
+
+FAULTY_ENTRY_ID = "FAULTY"
+
+
+def side_effect_async_get_data(x):
+    """Side effect on calling async get data which fills the connections attribute."""
+    x.connections = [
+        {
+            "departure": "DEPARTURE_0",
+            "number": 0,
+            "platform": 0,
+            "transfers": 0,
+            "duration": "10",
+            "delay": 0,
+        },
+        {
+            "departure": "DEPARTURE_1",
+            "number": 1,
+            "platform": 1,
+            "transfers": 0,
+            "duration": "10",
+            "delay": 0,
+        },
+        {
+            "departure": "DEPARTURE_2",
+            "number": 2,
+            "platform": 2,
+            "transfers": 0,
+            "duration": "10",
+            "delay": 0,
+        },
+    ]
+
+
+async def test_migration_1_2(hass: HomeAssistant) -> None:
+    """Test successful setup."""
+
+    with patch(
+        "homeassistant.components.swiss_public_transport.OpendataTransport.async_get_data",
+        autospec=True,
+        return_value=True,
+        side_effect=side_effect_async_get_data,
+    ):
+        # Set up mock objects
+        entity_registry = er.async_get(hass)
+
+        config_entry_faulty = MockConfigEntry(
+            domain=DOMAIN,
+            entry_id=FAULTY_ENTRY_ID,
+            data=MOCK_DATA_STEP,
+            title="MIGRATION_TEST",
+            unique_id=None,
+        )
+        config_entry_faulty.add_to_hass(hass)
+
+        # Setup the config entry
+        await async_setup_entry(hass, config_entry_faulty)
+        await hass.async_block_till_done()
+        # print(entity_registry.entities.values())
+        # print(entity_registry.entities._index)
+        assert entity_registry.async_is_registered(
+            entity_registry.entities.get_entity_id(
+                (Platform.SENSOR, DOMAIN, "None_departure")
+            )
+        )
+
+        # Migration: add unique_id and save new config entry, remove faulty entity from registry
+        assert await async_migrate_entry(hass, config_entry_faulty)
+        await hass.async_block_till_done()
+
+        # Check change in config entry
+        assert config_entry_faulty.version == 2
+        assert config_entry_faulty.unique_id == "test_start test_destination"
+
+        # Check faulty entity has been removed
+        assert not entity_registry.async_is_registered(
+            entity_registry.entities.get_entity_id(
+                (Platform.SENSOR, DOMAIN, "None_departure")
+            )
+        )
+        # Check migrated entity has been added
+        # print(config_entry_faulty.entry_id)
+        # print(entity_registry.entities.values())
+        # print(entity_registry._entities_data.keys())
+        assert entity_registry.async_is_registered(
+            entity_registry.entities.get_entity_id(
+                (Platform.SENSOR, DOMAIN, "test_start test_destination_departure")
+            )
+        )