diff --git a/homeassistant/components/drop_connect/__init__.py b/homeassistant/components/drop_connect/__init__.py index bc700456398..52b8f5a7d6e 100644 --- a/homeassistant/components/drop_connect/__init__.py +++ b/homeassistant/components/drop_connect/__init__.py @@ -7,12 +7,11 @@ from typing import TYPE_CHECKING from homeassistant.components import mqtt from homeassistant.components.mqtt import ReceiveMessage -from homeassistant.config_entries import ConfigEntry from homeassistant.const import Platform from homeassistant.core import HomeAssistant, callback -from .const import CONF_DATA_TOPIC, CONF_DEVICE_TYPE, DOMAIN -from .coordinator import DROPDeviceDataUpdateCoordinator +from .const import CONF_DATA_TOPIC, CONF_DEVICE_TYPE +from .coordinator import DROPConfigEntry, DROPDeviceDataUpdateCoordinator _LOGGER = logging.getLogger(__name__) @@ -24,7 +23,7 @@ PLATFORMS: list[Platform] = [ ] -async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool: +async def async_setup_entry(hass: HomeAssistant, config_entry: DROPConfigEntry) -> bool: """Set up DROP from a config entry.""" # Make sure MQTT integration is enabled and the client is available. @@ -34,9 +33,7 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b if TYPE_CHECKING: assert config_entry.unique_id is not None - drop_data_coordinator = DROPDeviceDataUpdateCoordinator( - hass, config_entry.unique_id - ) + drop_data_coordinator = DROPDeviceDataUpdateCoordinator(hass, config_entry) @callback def mqtt_callback(msg: ReceiveMessage) -> None: @@ -58,15 +55,13 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b config_entry.data[CONF_DATA_TOPIC], ) - hass.data.setdefault(DOMAIN, {})[config_entry.entry_id] = drop_data_coordinator + config_entry.runtime_data = drop_data_coordinator await hass.config_entries.async_forward_entry_setups(config_entry, PLATFORMS) return True -async def async_unload_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool: +async def async_unload_entry( + hass: HomeAssistant, config_entry: DROPConfigEntry +) -> bool: """Unload a config entry.""" - if unload_ok := await hass.config_entries.async_unload_platforms( - config_entry, PLATFORMS - ): - hass.data[DOMAIN].pop(config_entry.entry_id) - return unload_ok + return await hass.config_entries.async_unload_platforms(config_entry, PLATFORMS) diff --git a/homeassistant/components/drop_connect/binary_sensor.py b/homeassistant/components/drop_connect/binary_sensor.py index 093c5bcbb8e..bc8cf900610 100644 --- a/homeassistant/components/drop_connect/binary_sensor.py +++ b/homeassistant/components/drop_connect/binary_sensor.py @@ -11,7 +11,6 @@ from homeassistant.components.binary_sensor import ( BinarySensorEntity, BinarySensorEntityDescription, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback @@ -25,9 +24,8 @@ from .const import ( DEV_RO_FILTER, DEV_SALT_SENSOR, DEV_SOFTENER, - DOMAIN, ) -from .coordinator import DROPDeviceDataUpdateCoordinator +from .coordinator import DROPConfigEntry, DROPDeviceDataUpdateCoordinator from .entity import DROPEntity _LOGGER = logging.getLogger(__name__) @@ -106,7 +104,7 @@ DEVICE_BINARY_SENSORS: dict[str, list[str]] = { async def async_setup_entry( hass: HomeAssistant, - config_entry: ConfigEntry, + config_entry: DROPConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up the DROP binary sensors from config entry.""" @@ -116,9 +114,10 @@ async def async_setup_entry( config_entry.entry_id, ) + coordinator = config_entry.runtime_data if config_entry.data[CONF_DEVICE_TYPE] in DEVICE_BINARY_SENSORS: async_add_entities( - DROPBinarySensor(hass.data[DOMAIN][config_entry.entry_id], sensor) + DROPBinarySensor(coordinator, sensor) for sensor in BINARY_SENSORS if sensor.key in DEVICE_BINARY_SENSORS[config_entry.data[CONF_DEVICE_TYPE]] ) diff --git a/homeassistant/components/drop_connect/coordinator.py b/homeassistant/components/drop_connect/coordinator.py index 0861e091153..d37127d89ed 100644 --- a/homeassistant/components/drop_connect/coordinator.py +++ b/homeassistant/components/drop_connect/coordinator.py @@ -16,14 +16,19 @@ from .const import CONF_COMMAND_TOPIC, DOMAIN _LOGGER = logging.getLogger(__name__) +type DROPConfigEntry = ConfigEntry[DROPDeviceDataUpdateCoordinator] + + class DROPDeviceDataUpdateCoordinator(DataUpdateCoordinator[None]): """DROP device object.""" - config_entry: ConfigEntry + config_entry: DROPConfigEntry - def __init__(self, hass: HomeAssistant, unique_id: str) -> None: + def __init__(self, hass: HomeAssistant, entry: DROPConfigEntry) -> None: """Initialize the device.""" - super().__init__(hass, _LOGGER, name=f"{DOMAIN}-{unique_id}") + super().__init__( + hass, _LOGGER, config_entry=entry, name=f"{DOMAIN}-{entry.unique_id}" + ) self.drop_api = DropAPI() async def set_water(self, value: int) -> None: diff --git a/homeassistant/components/drop_connect/select.py b/homeassistant/components/drop_connect/select.py index ad06576c9f3..9e4c74b67e6 100644 --- a/homeassistant/components/drop_connect/select.py +++ b/homeassistant/components/drop_connect/select.py @@ -8,12 +8,11 @@ import logging from typing import Any from homeassistant.components.select import SelectEntity, SelectEntityDescription -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import CONF_DEVICE_TYPE, DEV_HUB, DOMAIN -from .coordinator import DROPDeviceDataUpdateCoordinator +from .const import CONF_DEVICE_TYPE, DEV_HUB +from .coordinator import DROPConfigEntry, DROPDeviceDataUpdateCoordinator from .entity import DROPEntity _LOGGER = logging.getLogger(__name__) @@ -50,7 +49,7 @@ DEVICE_SELECTS: dict[str, list[str]] = { async def async_setup_entry( hass: HomeAssistant, - config_entry: ConfigEntry, + config_entry: DROPConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up the DROP selects from config entry.""" @@ -60,9 +59,10 @@ async def async_setup_entry( config_entry.entry_id, ) + coordinator = config_entry.runtime_data if config_entry.data[CONF_DEVICE_TYPE] in DEVICE_SELECTS: async_add_entities( - DROPSelect(hass.data[DOMAIN][config_entry.entry_id], select) + DROPSelect(coordinator, select) for select in SELECTS if select.key in DEVICE_SELECTS[config_entry.data[CONF_DEVICE_TYPE]] ) diff --git a/homeassistant/components/drop_connect/sensor.py b/homeassistant/components/drop_connect/sensor.py index ad123ee13c7..5ec47ed9eb1 100644 --- a/homeassistant/components/drop_connect/sensor.py +++ b/homeassistant/components/drop_connect/sensor.py @@ -12,7 +12,6 @@ from homeassistant.components.sensor import ( SensorEntityDescription, SensorStateClass, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import ( CONCENTRATION_PARTS_PER_MILLION, PERCENTAGE, @@ -35,9 +34,8 @@ from .const import ( DEV_PUMP_CONTROLLER, DEV_RO_FILTER, DEV_SOFTENER, - DOMAIN, ) -from .coordinator import DROPDeviceDataUpdateCoordinator +from .coordinator import DROPConfigEntry, DROPDeviceDataUpdateCoordinator from .entity import DROPEntity _LOGGER = logging.getLogger(__name__) @@ -243,7 +241,7 @@ DEVICE_SENSORS: dict[str, list[str]] = { async def async_setup_entry( hass: HomeAssistant, - config_entry: ConfigEntry, + config_entry: DROPConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up the DROP sensors from config entry.""" @@ -253,9 +251,10 @@ async def async_setup_entry( config_entry.entry_id, ) + coordinator = config_entry.runtime_data if config_entry.data[CONF_DEVICE_TYPE] in DEVICE_SENSORS: async_add_entities( - DROPSensor(hass.data[DOMAIN][config_entry.entry_id], sensor) + DROPSensor(coordinator, sensor) for sensor in SENSORS if sensor.key in DEVICE_SENSORS[config_entry.data[CONF_DEVICE_TYPE]] ) diff --git a/homeassistant/components/drop_connect/switch.py b/homeassistant/components/drop_connect/switch.py index 98841d7ca24..404059d3196 100644 --- a/homeassistant/components/drop_connect/switch.py +++ b/homeassistant/components/drop_connect/switch.py @@ -8,7 +8,6 @@ import logging from typing import Any from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback @@ -18,9 +17,8 @@ from .const import ( DEV_HUB, DEV_PROTECTION_VALVE, DEV_SOFTENER, - DOMAIN, ) -from .coordinator import DROPDeviceDataUpdateCoordinator +from .coordinator import DROPConfigEntry, DROPDeviceDataUpdateCoordinator from .entity import DROPEntity _LOGGER = logging.getLogger(__name__) @@ -66,7 +64,7 @@ DEVICE_SWITCHES: dict[str, list[str]] = { async def async_setup_entry( hass: HomeAssistant, - config_entry: ConfigEntry, + config_entry: DROPConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up the DROP switches from config entry.""" @@ -76,9 +74,10 @@ async def async_setup_entry( config_entry.entry_id, ) + coordinator = config_entry.runtime_data if config_entry.data[CONF_DEVICE_TYPE] in DEVICE_SWITCHES: async_add_entities( - DROPSwitch(hass.data[DOMAIN][config_entry.entry_id], switch) + DROPSwitch(coordinator, switch) for switch in SWITCHES if switch.key in DEVICE_SWITCHES[config_entry.data[CONF_DEVICE_TYPE]] )