diff --git a/homeassistant/components/atag/__init__.py b/homeassistant/components/atag/__init__.py index 6ce647f7402..89f95f77870 100644 --- a/homeassistant/components/atag/__init__.py +++ b/homeassistant/components/atag/__init__.py @@ -1,22 +1,21 @@ """The ATAG Integration.""" -from homeassistant.config_entries import ConfigEntry from homeassistant.const import Platform from homeassistant.core import HomeAssistant -from .coordinator import AtagDataUpdateCoordinator +from .coordinator import AtagConfigEntry, AtagDataUpdateCoordinator DOMAIN = "atag" PLATFORMS = [Platform.CLIMATE, Platform.SENSOR, Platform.WATER_HEATER] -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_setup_entry(hass: HomeAssistant, entry: AtagConfigEntry) -> bool: """Set up Atag integration from a config entry.""" coordinator = AtagDataUpdateCoordinator(hass, entry) await coordinator.async_config_entry_first_refresh() - hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator + entry.runtime_data = coordinator if entry.unique_id is None: hass.config_entries.async_update_entry(entry, unique_id=coordinator.atag.id) @@ -25,10 +24,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: return True -async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_unload_entry(hass: HomeAssistant, entry: AtagConfigEntry) -> bool: """Unload Atag config entry.""" - unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS) - - if unload_ok: - hass.data[DOMAIN].pop(entry.entry_id) - return unload_ok + return await hass.config_entries.async_unload_platforms(entry, PLATFORMS) diff --git a/homeassistant/components/atag/climate.py b/homeassistant/components/atag/climate.py index 744339e6e23..daeb64f7f0a 100644 --- a/homeassistant/components/atag/climate.py +++ b/homeassistant/components/atag/climate.py @@ -12,13 +12,12 @@ from homeassistant.components.climate import ( HVACAction, HVACMode, ) -from homeassistant.config_entries import ConfigEntry -from homeassistant.const import ATTR_TEMPERATURE, Platform +from homeassistant.const import ATTR_TEMPERATURE from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.util.enum import try_parse_enum -from . import DOMAIN +from .coordinator import AtagConfigEntry, AtagDataUpdateCoordinator from .entity import AtagEntity PRESET_MAP = { @@ -33,11 +32,10 @@ HVAC_MODES = [HVACMode.AUTO, HVACMode.HEAT] async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, entry: AtagConfigEntry, async_add_entities: AddEntitiesCallback ) -> None: """Load a config entry.""" - coordinator = hass.data[DOMAIN][entry.entry_id] - async_add_entities([AtagThermostat(coordinator, Platform.CLIMATE)]) + async_add_entities([AtagThermostat(entry.runtime_data, "climate")]) class AtagThermostat(AtagEntity, ClimateEntity): @@ -50,7 +48,7 @@ class AtagThermostat(AtagEntity, ClimateEntity): ) _enable_turn_on_off_backwards_compatibility = False - def __init__(self, coordinator, atag_id): + def __init__(self, coordinator: AtagDataUpdateCoordinator, atag_id: str) -> None: """Initialize an Atag climate device.""" super().__init__(coordinator, atag_id) self._attr_temperature_unit = coordinator.atag.climate.temp_unit diff --git a/homeassistant/components/atag/coordinator.py b/homeassistant/components/atag/coordinator.py index 84f8f3eb05e..6d542471384 100644 --- a/homeassistant/components/atag/coordinator.py +++ b/homeassistant/components/atag/coordinator.py @@ -13,6 +13,8 @@ from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, Upda _LOGGER = logging.getLogger(__name__) +type AtagConfigEntry = ConfigEntry[AtagDataUpdateCoordinator] + class AtagDataUpdateCoordinator(DataUpdateCoordinator[None]): """Atag data update coordinator.""" diff --git a/homeassistant/components/atag/sensor.py b/homeassistant/components/atag/sensor.py index 85fa403f3a9..bd39f0b3458 100644 --- a/homeassistant/components/atag/sensor.py +++ b/homeassistant/components/atag/sensor.py @@ -1,7 +1,6 @@ """Initialization of ATAG One sensor platform.""" from homeassistant.components.sensor import SensorDeviceClass, SensorEntity -from homeassistant.config_entries import ConfigEntry from homeassistant.const import ( PERCENTAGE, UnitOfPressure, @@ -11,7 +10,7 @@ from homeassistant.const import ( from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback -from . import DOMAIN +from .coordinator import AtagConfigEntry, AtagDataUpdateCoordinator from .entity import AtagEntity SENSORS = { @@ -28,18 +27,18 @@ SENSORS = { async def async_setup_entry( hass: HomeAssistant, - config_entry: ConfigEntry, + config_entry: AtagConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Initialize sensor platform from config entry.""" - coordinator = hass.data[DOMAIN][config_entry.entry_id] + coordinator = config_entry.runtime_data async_add_entities([AtagSensor(coordinator, sensor) for sensor in SENSORS]) class AtagSensor(AtagEntity, SensorEntity): """Representation of a AtagOne Sensor.""" - def __init__(self, coordinator, sensor): + def __init__(self, coordinator: AtagDataUpdateCoordinator, sensor: str) -> None: """Initialize Atag sensor.""" super().__init__(coordinator, SENSORS[sensor]) self._attr_name = sensor diff --git a/homeassistant/components/atag/water_heater.py b/homeassistant/components/atag/water_heater.py index e95d7e03386..6b013b36885 100644 --- a/homeassistant/components/atag/water_heater.py +++ b/homeassistant/components/atag/water_heater.py @@ -7,12 +7,11 @@ from homeassistant.components.water_heater import ( STATE_PERFORMANCE, WaterHeaterEntity, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import ATTR_TEMPERATURE, STATE_OFF, Platform, UnitOfTemperature from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback -from . import DOMAIN +from .coordinator import AtagConfigEntry from .entity import AtagEntity OPERATION_LIST = [STATE_OFF, STATE_ECO, STATE_PERFORMANCE] @@ -20,12 +19,13 @@ OPERATION_LIST = [STATE_OFF, STATE_ECO, STATE_PERFORMANCE] async def async_setup_entry( hass: HomeAssistant, - config_entry: ConfigEntry, + config_entry: AtagConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Initialize DHW device from config entry.""" - coordinator = hass.data[DOMAIN][config_entry.entry_id] - async_add_entities([AtagWaterHeater(coordinator, Platform.WATER_HEATER)]) + async_add_entities( + [AtagWaterHeater(config_entry.runtime_data, Platform.WATER_HEATER)] + ) class AtagWaterHeater(AtagEntity, WaterHeaterEntity): diff --git a/tests/components/atag/test_climate.py b/tests/components/atag/test_climate.py index cfc0c2a5d47..b4f2a0f3f0f 100644 --- a/tests/components/atag/test_climate.py +++ b/tests/components/atag/test_climate.py @@ -2,7 +2,8 @@ from unittest.mock import PropertyMock, patch -from homeassistant.components.atag.climate import DOMAIN, PRESET_MAP +from homeassistant.components.atag import DOMAIN +from homeassistant.components.atag.climate import PRESET_MAP from homeassistant.components.climate import ( ATTR_HVAC_ACTION, ATTR_HVAC_MODE, @@ -104,7 +105,7 @@ async def test_update_failed( entry = await init_integration(hass, aioclient_mock) await async_setup_component(hass, HA_DOMAIN, {}) assert hass.states.get(CLIMATE_ID).state == HVACMode.HEAT - coordinator = hass.data[DOMAIN][entry.entry_id] + coordinator = entry.runtime_data with patch("pyatag.AtagOne.update", side_effect=TimeoutError) as updater: await coordinator.async_refresh() await hass.async_block_till_done() diff --git a/tests/components/atag/test_init.py b/tests/components/atag/test_init.py index 59f38ae7bfe..7c65150fbf6 100644 --- a/tests/components/atag/test_init.py +++ b/tests/components/atag/test_init.py @@ -1,6 +1,5 @@ """Tests for the ATAG integration.""" -from homeassistant.components.atag import DOMAIN from homeassistant.config_entries import ConfigEntryState from homeassistant.core import HomeAssistant @@ -23,7 +22,7 @@ async def test_unload_config_entry( ) -> None: """Test the ATAG configuration entry unloading.""" entry = await init_integration(hass, aioclient_mock) - assert hass.data[DOMAIN] + assert entry.runtime_data await hass.config_entries.async_unload(entry.entry_id) await hass.async_block_till_done() - assert not hass.data.get(DOMAIN) + assert not hasattr(entry, "runtime_data")