diff --git a/homeassistant/components/swiss_public_transport/__init__.py b/homeassistant/components/swiss_public_transport/__init__.py index e0c3f61a85e15a..f3e486bcfe90e1 100644 --- a/homeassistant/components/swiss_public_transport/__init__.py +++ b/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,9 +80,11 @@ 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 @@ -89,13 +92,18 @@ async def async_migrate_entry( 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) diff --git a/tests/components/swiss_public_transport/test_migration.py b/tests/components/swiss_public_transport/test_migration.py new file mode 100644 index 00000000000000..ffe89002bb9ae6 --- /dev/null +++ b/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") + ) + )