Make sure RainMachine data storage conforms to standards (#57816)

This commit is contained in:
Aaron Bach 2021-10-22 04:14:58 -06:00 committed by GitHub
parent 8dfa628af0
commit 8bc10db0bb
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 22 additions and 21 deletions

View File

@ -4,7 +4,7 @@ from __future__ import annotations
import asyncio import asyncio
from datetime import timedelta from datetime import timedelta
from functools import partial from functools import partial
from typing import Any from typing import Any, cast
from regenmaschine import Client from regenmaschine import Client
from regenmaschine.controller import Controller from regenmaschine.controller import Controller
@ -123,14 +123,13 @@ def async_get_controller_for_service_call(
hass: HomeAssistant, call: ServiceCall hass: HomeAssistant, call: ServiceCall
) -> Controller: ) -> Controller:
"""Get the controller related to a service call (by device ID).""" """Get the controller related to a service call (by device ID)."""
controllers: dict[str, Controller] = hass.data[DOMAIN][DATA_CONTROLLER]
device_id = call.data[CONF_DEVICE_ID] device_id = call.data[CONF_DEVICE_ID]
device_registry = dr.async_get(hass) device_registry = dr.async_get(hass)
if device_entry := device_registry.async_get(device_id): if device_entry := device_registry.async_get(device_id):
for entry_id in device_entry.config_entries: for entry_id in device_entry.config_entries:
if entry_id in controllers: if controller := hass.data[DOMAIN][entry_id][DATA_CONTROLLER]:
return controllers[entry_id] return cast(Controller, controller)
raise ValueError(f"No controller for device ID: {device_id}") raise ValueError(f"No controller for device ID: {device_id}")
@ -145,10 +144,10 @@ async def async_update_programs_and_zones(
""" """
await asyncio.gather( await asyncio.gather(
*[ *[
hass.data[DOMAIN][DATA_COORDINATOR][entry.entry_id][ hass.data[DOMAIN][entry.entry_id][DATA_COORDINATOR][
DATA_PROGRAMS DATA_PROGRAMS
].async_refresh(), ].async_refresh(),
hass.data[DOMAIN][DATA_COORDINATOR][entry.entry_id][ hass.data[DOMAIN][entry.entry_id][DATA_COORDINATOR][
DATA_ZONES DATA_ZONES
].async_refresh(), ].async_refresh(),
] ]
@ -157,8 +156,9 @@ async def async_update_programs_and_zones(
async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
"""Set up RainMachine as config entry.""" """Set up RainMachine as config entry."""
hass.data.setdefault(DOMAIN, {DATA_CONTROLLER: {}, DATA_COORDINATOR: {}}) hass.data.setdefault(DOMAIN, {})
hass.data[DOMAIN][DATA_COORDINATOR][entry.entry_id] = {} hass.data[DOMAIN][entry.entry_id] = {}
websession = aiohttp_client.async_get_clientsession(hass) websession = aiohttp_client.async_get_clientsession(hass)
client = Client(session=websession) client = Client(session=websession)
@ -174,8 +174,8 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
# regenmaschine can load multiple controllers at once, but we only grab the one # regenmaschine can load multiple controllers at once, but we only grab the one
# we loaded above: # we loaded above:
controller = hass.data[DOMAIN][DATA_CONTROLLER][ controller = hass.data[DOMAIN][entry.entry_id][
entry.entry_id DATA_CONTROLLER
] = get_client_controller(client) ] = get_client_controller(client)
entry_updates: dict[str, Any] = {} entry_updates: dict[str, Any] = {}
@ -215,6 +215,8 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
return data return data
controller_init_tasks = [] controller_init_tasks = []
coordinators = {}
for api_category in ( for api_category in (
DATA_PROGRAMS, DATA_PROGRAMS,
DATA_PROVISION_SETTINGS, DATA_PROVISION_SETTINGS,
@ -222,9 +224,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
DATA_RESTRICTIONS_UNIVERSAL, DATA_RESTRICTIONS_UNIVERSAL,
DATA_ZONES, DATA_ZONES,
): ):
coordinator = hass.data[DOMAIN][DATA_COORDINATOR][entry.entry_id][ coordinator = coordinators[api_category] = DataUpdateCoordinator(
api_category
] = DataUpdateCoordinator(
hass, hass,
LOGGER, LOGGER,
name=f'{controller.name} ("{api_category}")', name=f'{controller.name} ("{api_category}")',
@ -234,6 +234,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
controller_init_tasks.append(coordinator.async_refresh()) controller_init_tasks.append(coordinator.async_refresh())
await asyncio.gather(*controller_init_tasks) await asyncio.gather(*controller_init_tasks)
hass.data[DOMAIN][entry.entry_id][DATA_COORDINATOR] = coordinators
hass.config_entries.async_setup_platforms(entry, PLATFORMS) hass.config_entries.async_setup_platforms(entry, PLATFORMS)
@ -297,7 +298,7 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
"""Unload an RainMachine config entry.""" """Unload an RainMachine config entry."""
unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS) unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
if unload_ok: if unload_ok:
hass.data[DOMAIN][DATA_COORDINATOR].pop(entry.entry_id) hass.data[DOMAIN].pop(entry.entry_id)
if len(hass.config_entries.async_entries(DOMAIN)) == 1: if len(hass.config_entries.async_entries(DOMAIN)) == 1:
# If this is the last instance of RainMachine, deregister any services defined # If this is the last instance of RainMachine, deregister any services defined

View File

@ -115,8 +115,8 @@ async def async_setup_entry(
hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
) -> None: ) -> None:
"""Set up RainMachine binary sensors based on a config entry.""" """Set up RainMachine binary sensors based on a config entry."""
controller = hass.data[DOMAIN][DATA_CONTROLLER][entry.entry_id] controller = hass.data[DOMAIN][entry.entry_id][DATA_CONTROLLER]
coordinators = hass.data[DOMAIN][DATA_COORDINATOR][entry.entry_id] coordinators = hass.data[DOMAIN][entry.entry_id][DATA_COORDINATOR]
@callback @callback
def async_get_sensor(api_category: str) -> partial: def async_get_sensor(api_category: str) -> partial:

View File

@ -101,8 +101,8 @@ async def async_setup_entry(
hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
) -> None: ) -> None:
"""Set up RainMachine sensors based on a config entry.""" """Set up RainMachine sensors based on a config entry."""
controller = hass.data[DOMAIN][DATA_CONTROLLER][entry.entry_id] controller = hass.data[DOMAIN][entry.entry_id][DATA_CONTROLLER]
coordinators = hass.data[DOMAIN][DATA_COORDINATOR][entry.entry_id] coordinators = hass.data[DOMAIN][entry.entry_id][DATA_COORDINATOR]
@callback @callback
def async_get_sensor(api_category: str) -> partial: def async_get_sensor(api_category: str) -> partial:

View File

@ -149,11 +149,11 @@ async def async_setup_entry(
): ):
platform.async_register_entity_service(service_name, schema, method) platform.async_register_entity_service(service_name, schema, method)
controller = hass.data[DOMAIN][DATA_CONTROLLER][entry.entry_id] controller = hass.data[DOMAIN][entry.entry_id][DATA_CONTROLLER]
programs_coordinator = hass.data[DOMAIN][DATA_COORDINATOR][entry.entry_id][ programs_coordinator = hass.data[DOMAIN][entry.entry_id][DATA_COORDINATOR][
DATA_PROGRAMS DATA_PROGRAMS
] ]
zones_coordinator = hass.data[DOMAIN][DATA_COORDINATOR][entry.entry_id][DATA_ZONES] zones_coordinator = hass.data[DOMAIN][entry.entry_id][DATA_COORDINATOR][DATA_ZONES]
entities: list[RainMachineProgram | RainMachineZone] = [ entities: list[RainMachineProgram | RainMachineZone] = [
RainMachineProgram( RainMachineProgram(