From b2dbaf7feb6f92bae1f4d22d2de8446ca2011b64 Mon Sep 17 00:00:00 2001
From: miaucl <cyrill.raccaud@gmail.com>
Date: Thu, 4 Jan 2024 10:34:31 +0000
Subject: [PATCH] bring

---
 .../swiss_public_transport/__init__.py        | 36 +++++++++++++++++++
 .../swiss_public_transport/config_flow.py     |  8 ++++-
 2 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/homeassistant/components/swiss_public_transport/__init__.py b/homeassistant/components/swiss_public_transport/__init__.py
index 9e01a07416f912..183d337fb5414d 100644
--- a/homeassistant/components/swiss_public_transport/__init__.py
+++ b/homeassistant/components/swiss_public_transport/__init__.py
@@ -10,6 +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.aiohttp_client import async_get_clientsession
 
 from .const import CONF_DESTINATION, CONF_START, DOMAIN
@@ -65,3 +66,38 @@ async def async_unload_entry(
         hass.data[DOMAIN].pop(entry.entry_id)
 
     return unload_ok
+
+
+async def async_migrate_entry(
+    hass: core.HomeAssistant, config_entry: config_entries.ConfigEntry
+) -> bool:
+    """Migrate config entry."""
+    _LOGGER.info(config_entry.version)
+    _LOGGER.debug("Migrating from version %s", config_entry.version)
+
+    if config_entry.version > 3:
+        # This means the user has downgraded from a future version
+        return False
+
+    if config_entry.version == 1:
+        new = {**config_entry.data}
+
+        # Remove wrongly registered devices
+        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]}"
+        )
+
+        config_entry.version = 2
+        hass.config_entries.async_update_entry(config_entry, data=new)
+
+    _LOGGER.debug("Migration to version %s successful", config_entry.version)
+
+    return True
diff --git a/homeassistant/components/swiss_public_transport/config_flow.py b/homeassistant/components/swiss_public_transport/config_flow.py
index 63eca1efe96a17..fb3c135a56fcda 100644
--- a/homeassistant/components/swiss_public_transport/config_flow.py
+++ b/homeassistant/components/swiss_public_transport/config_flow.py
@@ -30,7 +30,7 @@
 class SwissPublicTransportConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
     """Swiss public transport config flow."""
 
-    VERSION = 1
+    VERSION = 2
 
     async def async_step_user(
         self, user_input: dict[str, Any] | None = None
@@ -59,6 +59,9 @@ async def async_step_user(
                 _LOGGER.exception("Unknown error")
                 errors["base"] = "unknown"
             else:
+                await self.async_set_unique_id(
+                    f"{user_input[CONF_START]} {user_input[CONF_DESTINATION]}"
+                )
                 return self.async_create_entry(
                     title=f"{user_input[CONF_START]} {user_input[CONF_DESTINATION]}",
                     data=user_input,
@@ -98,6 +101,9 @@ async def async_step_import(self, import_input: dict[str, Any]) -> FlowResult:
             )
             return self.async_abort(reason="unknown")
 
+        await self.async_set_unique_id(
+            f"{import_input[CONF_START]} {import_input[CONF_DESTINATION]}"
+        )
         return self.async_create_entry(
             title=import_input[CONF_NAME],
             data=import_input,