Use runtime_data in drop_connect (#136442)

This commit is contained in:
epenet 2025-01-28 14:10:09 +01:00 committed by GitHub
parent 2c3cd6e119
commit 9897e4d3e4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 34 additions and 37 deletions

View File

@ -7,12 +7,11 @@ from typing import TYPE_CHECKING
from homeassistant.components import mqtt from homeassistant.components import mqtt
from homeassistant.components.mqtt import ReceiveMessage from homeassistant.components.mqtt import ReceiveMessage
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import Platform from homeassistant.const import Platform
from homeassistant.core import HomeAssistant, callback from homeassistant.core import HomeAssistant, callback
from .const import CONF_DATA_TOPIC, CONF_DEVICE_TYPE, DOMAIN from .const import CONF_DATA_TOPIC, CONF_DEVICE_TYPE
from .coordinator import DROPDeviceDataUpdateCoordinator from .coordinator import DROPConfigEntry, DROPDeviceDataUpdateCoordinator
_LOGGER = logging.getLogger(__name__) _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.""" """Set up DROP from a config entry."""
# Make sure MQTT integration is enabled and the client is available. # 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: if TYPE_CHECKING:
assert config_entry.unique_id is not None assert config_entry.unique_id is not None
drop_data_coordinator = DROPDeviceDataUpdateCoordinator( drop_data_coordinator = DROPDeviceDataUpdateCoordinator(hass, config_entry)
hass, config_entry.unique_id
)
@callback @callback
def mqtt_callback(msg: ReceiveMessage) -> None: 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], 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) await hass.config_entries.async_forward_entry_setups(config_entry, PLATFORMS)
return True 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.""" """Unload a config entry."""
if unload_ok := await hass.config_entries.async_unload_platforms( return await hass.config_entries.async_unload_platforms(config_entry, PLATFORMS)
config_entry, PLATFORMS
):
hass.data[DOMAIN].pop(config_entry.entry_id)
return unload_ok

View File

@ -11,7 +11,6 @@ from homeassistant.components.binary_sensor import (
BinarySensorEntity, BinarySensorEntity,
BinarySensorEntityDescription, BinarySensorEntityDescription,
) )
from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant from homeassistant.core import HomeAssistant
from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.entity_platform import AddEntitiesCallback
@ -25,9 +24,8 @@ from .const import (
DEV_RO_FILTER, DEV_RO_FILTER,
DEV_SALT_SENSOR, DEV_SALT_SENSOR,
DEV_SOFTENER, DEV_SOFTENER,
DOMAIN,
) )
from .coordinator import DROPDeviceDataUpdateCoordinator from .coordinator import DROPConfigEntry, DROPDeviceDataUpdateCoordinator
from .entity import DROPEntity from .entity import DROPEntity
_LOGGER = logging.getLogger(__name__) _LOGGER = logging.getLogger(__name__)
@ -106,7 +104,7 @@ DEVICE_BINARY_SENSORS: dict[str, list[str]] = {
async def async_setup_entry( async def async_setup_entry(
hass: HomeAssistant, hass: HomeAssistant,
config_entry: ConfigEntry, config_entry: DROPConfigEntry,
async_add_entities: AddEntitiesCallback, async_add_entities: AddEntitiesCallback,
) -> None: ) -> None:
"""Set up the DROP binary sensors from config entry.""" """Set up the DROP binary sensors from config entry."""
@ -116,9 +114,10 @@ async def async_setup_entry(
config_entry.entry_id, config_entry.entry_id,
) )
coordinator = config_entry.runtime_data
if config_entry.data[CONF_DEVICE_TYPE] in DEVICE_BINARY_SENSORS: if config_entry.data[CONF_DEVICE_TYPE] in DEVICE_BINARY_SENSORS:
async_add_entities( async_add_entities(
DROPBinarySensor(hass.data[DOMAIN][config_entry.entry_id], sensor) DROPBinarySensor(coordinator, sensor)
for sensor in BINARY_SENSORS for sensor in BINARY_SENSORS
if sensor.key in DEVICE_BINARY_SENSORS[config_entry.data[CONF_DEVICE_TYPE]] if sensor.key in DEVICE_BINARY_SENSORS[config_entry.data[CONF_DEVICE_TYPE]]
) )

View File

@ -16,14 +16,19 @@ from .const import CONF_COMMAND_TOPIC, DOMAIN
_LOGGER = logging.getLogger(__name__) _LOGGER = logging.getLogger(__name__)
type DROPConfigEntry = ConfigEntry[DROPDeviceDataUpdateCoordinator]
class DROPDeviceDataUpdateCoordinator(DataUpdateCoordinator[None]): class DROPDeviceDataUpdateCoordinator(DataUpdateCoordinator[None]):
"""DROP device object.""" """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.""" """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() self.drop_api = DropAPI()
async def set_water(self, value: int) -> None: async def set_water(self, value: int) -> None:

View File

@ -8,12 +8,11 @@ import logging
from typing import Any from typing import Any
from homeassistant.components.select import SelectEntity, SelectEntityDescription from homeassistant.components.select import SelectEntity, SelectEntityDescription
from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant from homeassistant.core import HomeAssistant
from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.entity_platform import AddEntitiesCallback
from .const import CONF_DEVICE_TYPE, DEV_HUB, DOMAIN from .const import CONF_DEVICE_TYPE, DEV_HUB
from .coordinator import DROPDeviceDataUpdateCoordinator from .coordinator import DROPConfigEntry, DROPDeviceDataUpdateCoordinator
from .entity import DROPEntity from .entity import DROPEntity
_LOGGER = logging.getLogger(__name__) _LOGGER = logging.getLogger(__name__)
@ -50,7 +49,7 @@ DEVICE_SELECTS: dict[str, list[str]] = {
async def async_setup_entry( async def async_setup_entry(
hass: HomeAssistant, hass: HomeAssistant,
config_entry: ConfigEntry, config_entry: DROPConfigEntry,
async_add_entities: AddEntitiesCallback, async_add_entities: AddEntitiesCallback,
) -> None: ) -> None:
"""Set up the DROP selects from config entry.""" """Set up the DROP selects from config entry."""
@ -60,9 +59,10 @@ async def async_setup_entry(
config_entry.entry_id, config_entry.entry_id,
) )
coordinator = config_entry.runtime_data
if config_entry.data[CONF_DEVICE_TYPE] in DEVICE_SELECTS: if config_entry.data[CONF_DEVICE_TYPE] in DEVICE_SELECTS:
async_add_entities( async_add_entities(
DROPSelect(hass.data[DOMAIN][config_entry.entry_id], select) DROPSelect(coordinator, select)
for select in SELECTS for select in SELECTS
if select.key in DEVICE_SELECTS[config_entry.data[CONF_DEVICE_TYPE]] if select.key in DEVICE_SELECTS[config_entry.data[CONF_DEVICE_TYPE]]
) )

View File

@ -12,7 +12,6 @@ from homeassistant.components.sensor import (
SensorEntityDescription, SensorEntityDescription,
SensorStateClass, SensorStateClass,
) )
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import ( from homeassistant.const import (
CONCENTRATION_PARTS_PER_MILLION, CONCENTRATION_PARTS_PER_MILLION,
PERCENTAGE, PERCENTAGE,
@ -35,9 +34,8 @@ from .const import (
DEV_PUMP_CONTROLLER, DEV_PUMP_CONTROLLER,
DEV_RO_FILTER, DEV_RO_FILTER,
DEV_SOFTENER, DEV_SOFTENER,
DOMAIN,
) )
from .coordinator import DROPDeviceDataUpdateCoordinator from .coordinator import DROPConfigEntry, DROPDeviceDataUpdateCoordinator
from .entity import DROPEntity from .entity import DROPEntity
_LOGGER = logging.getLogger(__name__) _LOGGER = logging.getLogger(__name__)
@ -243,7 +241,7 @@ DEVICE_SENSORS: dict[str, list[str]] = {
async def async_setup_entry( async def async_setup_entry(
hass: HomeAssistant, hass: HomeAssistant,
config_entry: ConfigEntry, config_entry: DROPConfigEntry,
async_add_entities: AddEntitiesCallback, async_add_entities: AddEntitiesCallback,
) -> None: ) -> None:
"""Set up the DROP sensors from config entry.""" """Set up the DROP sensors from config entry."""
@ -253,9 +251,10 @@ async def async_setup_entry(
config_entry.entry_id, config_entry.entry_id,
) )
coordinator = config_entry.runtime_data
if config_entry.data[CONF_DEVICE_TYPE] in DEVICE_SENSORS: if config_entry.data[CONF_DEVICE_TYPE] in DEVICE_SENSORS:
async_add_entities( async_add_entities(
DROPSensor(hass.data[DOMAIN][config_entry.entry_id], sensor) DROPSensor(coordinator, sensor)
for sensor in SENSORS for sensor in SENSORS
if sensor.key in DEVICE_SENSORS[config_entry.data[CONF_DEVICE_TYPE]] if sensor.key in DEVICE_SENSORS[config_entry.data[CONF_DEVICE_TYPE]]
) )

View File

@ -8,7 +8,6 @@ import logging
from typing import Any from typing import Any
from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription
from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant from homeassistant.core import HomeAssistant
from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.entity_platform import AddEntitiesCallback
@ -18,9 +17,8 @@ from .const import (
DEV_HUB, DEV_HUB,
DEV_PROTECTION_VALVE, DEV_PROTECTION_VALVE,
DEV_SOFTENER, DEV_SOFTENER,
DOMAIN,
) )
from .coordinator import DROPDeviceDataUpdateCoordinator from .coordinator import DROPConfigEntry, DROPDeviceDataUpdateCoordinator
from .entity import DROPEntity from .entity import DROPEntity
_LOGGER = logging.getLogger(__name__) _LOGGER = logging.getLogger(__name__)
@ -66,7 +64,7 @@ DEVICE_SWITCHES: dict[str, list[str]] = {
async def async_setup_entry( async def async_setup_entry(
hass: HomeAssistant, hass: HomeAssistant,
config_entry: ConfigEntry, config_entry: DROPConfigEntry,
async_add_entities: AddEntitiesCallback, async_add_entities: AddEntitiesCallback,
) -> None: ) -> None:
"""Set up the DROP switches from config entry.""" """Set up the DROP switches from config entry."""
@ -76,9 +74,10 @@ async def async_setup_entry(
config_entry.entry_id, config_entry.entry_id,
) )
coordinator = config_entry.runtime_data
if config_entry.data[CONF_DEVICE_TYPE] in DEVICE_SWITCHES: if config_entry.data[CONF_DEVICE_TYPE] in DEVICE_SWITCHES:
async_add_entities( async_add_entities(
DROPSwitch(hass.data[DOMAIN][config_entry.entry_id], switch) DROPSwitch(coordinator, switch)
for switch in SWITCHES for switch in SWITCHES
if switch.key in DEVICE_SWITCHES[config_entry.data[CONF_DEVICE_TYPE]] if switch.key in DEVICE_SWITCHES[config_entry.data[CONF_DEVICE_TYPE]]
) )