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

View File

@ -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]]
)

View File

@ -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:

View File

@ -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]]
)

View File

@ -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]]
)

View File

@ -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]]
)