test_migration.py

"""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")
            )
        )