From 8bc10db0bb79f780922b1bddc51b3204fca58bf7 Mon Sep 17 00:00:00 2001 From: Aaron Bach Date: Fri, 22 Oct 2021 04:14:58 -0600 Subject: [PATCH] Make sure RainMachine data storage conforms to standards (#57816) --- .../components/rainmachine/__init__.py | 29 ++++++++++--------- .../components/rainmachine/binary_sensor.py | 4 +-- .../components/rainmachine/sensor.py | 4 +-- .../components/rainmachine/switch.py | 6 ++-- 4 files changed, 22 insertions(+), 21 deletions(-) diff --git a/homeassistant/components/rainmachine/__init__.py b/homeassistant/components/rainmachine/__init__.py index b8ea030cb71..2c2cb0ef3c1 100644 --- a/homeassistant/components/rainmachine/__init__.py +++ b/homeassistant/components/rainmachine/__init__.py @@ -4,7 +4,7 @@ from __future__ import annotations import asyncio from datetime import timedelta from functools import partial -from typing import Any +from typing import Any, cast from regenmaschine import Client from regenmaschine.controller import Controller @@ -123,14 +123,13 @@ def async_get_controller_for_service_call( hass: HomeAssistant, call: ServiceCall ) -> Controller: """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_registry = dr.async_get(hass) if device_entry := device_registry.async_get(device_id): for entry_id in device_entry.config_entries: - if entry_id in controllers: - return controllers[entry_id] + if controller := hass.data[DOMAIN][entry_id][DATA_CONTROLLER]: + return cast(Controller, controller) raise ValueError(f"No controller for device ID: {device_id}") @@ -145,10 +144,10 @@ async def async_update_programs_and_zones( """ await asyncio.gather( *[ - hass.data[DOMAIN][DATA_COORDINATOR][entry.entry_id][ + hass.data[DOMAIN][entry.entry_id][DATA_COORDINATOR][ DATA_PROGRAMS ].async_refresh(), - hass.data[DOMAIN][DATA_COORDINATOR][entry.entry_id][ + hass.data[DOMAIN][entry.entry_id][DATA_COORDINATOR][ DATA_ZONES ].async_refresh(), ] @@ -157,8 +156,9 @@ async def async_update_programs_and_zones( async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Set up RainMachine as config entry.""" - hass.data.setdefault(DOMAIN, {DATA_CONTROLLER: {}, DATA_COORDINATOR: {}}) - hass.data[DOMAIN][DATA_COORDINATOR][entry.entry_id] = {} + hass.data.setdefault(DOMAIN, {}) + hass.data[DOMAIN][entry.entry_id] = {} + websession = aiohttp_client.async_get_clientsession(hass) 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 # we loaded above: - controller = hass.data[DOMAIN][DATA_CONTROLLER][ - entry.entry_id + controller = hass.data[DOMAIN][entry.entry_id][ + DATA_CONTROLLER ] = get_client_controller(client) entry_updates: dict[str, Any] = {} @@ -215,6 +215,8 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: return data controller_init_tasks = [] + coordinators = {} + for api_category in ( DATA_PROGRAMS, DATA_PROVISION_SETTINGS, @@ -222,9 +224,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: DATA_RESTRICTIONS_UNIVERSAL, DATA_ZONES, ): - coordinator = hass.data[DOMAIN][DATA_COORDINATOR][entry.entry_id][ - api_category - ] = DataUpdateCoordinator( + coordinator = coordinators[api_category] = DataUpdateCoordinator( hass, LOGGER, 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()) await asyncio.gather(*controller_init_tasks) + hass.data[DOMAIN][entry.entry_id][DATA_COORDINATOR] = coordinators 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_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS) 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 this is the last instance of RainMachine, deregister any services defined diff --git a/homeassistant/components/rainmachine/binary_sensor.py b/homeassistant/components/rainmachine/binary_sensor.py index 55d009ad6f5..853b9f24b33 100644 --- a/homeassistant/components/rainmachine/binary_sensor.py +++ b/homeassistant/components/rainmachine/binary_sensor.py @@ -115,8 +115,8 @@ async def async_setup_entry( hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback ) -> None: """Set up RainMachine binary sensors based on a config entry.""" - controller = hass.data[DOMAIN][DATA_CONTROLLER][entry.entry_id] - coordinators = hass.data[DOMAIN][DATA_COORDINATOR][entry.entry_id] + controller = hass.data[DOMAIN][entry.entry_id][DATA_CONTROLLER] + coordinators = hass.data[DOMAIN][entry.entry_id][DATA_COORDINATOR] @callback def async_get_sensor(api_category: str) -> partial: diff --git a/homeassistant/components/rainmachine/sensor.py b/homeassistant/components/rainmachine/sensor.py index 6c95552e26f..dee0f1f6e57 100644 --- a/homeassistant/components/rainmachine/sensor.py +++ b/homeassistant/components/rainmachine/sensor.py @@ -101,8 +101,8 @@ async def async_setup_entry( hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback ) -> None: """Set up RainMachine sensors based on a config entry.""" - controller = hass.data[DOMAIN][DATA_CONTROLLER][entry.entry_id] - coordinators = hass.data[DOMAIN][DATA_COORDINATOR][entry.entry_id] + controller = hass.data[DOMAIN][entry.entry_id][DATA_CONTROLLER] + coordinators = hass.data[DOMAIN][entry.entry_id][DATA_COORDINATOR] @callback def async_get_sensor(api_category: str) -> partial: diff --git a/homeassistant/components/rainmachine/switch.py b/homeassistant/components/rainmachine/switch.py index e693664bab4..345da380316 100644 --- a/homeassistant/components/rainmachine/switch.py +++ b/homeassistant/components/rainmachine/switch.py @@ -149,11 +149,11 @@ async def async_setup_entry( ): platform.async_register_entity_service(service_name, schema, method) - controller = hass.data[DOMAIN][DATA_CONTROLLER][entry.entry_id] - programs_coordinator = hass.data[DOMAIN][DATA_COORDINATOR][entry.entry_id][ + controller = hass.data[DOMAIN][entry.entry_id][DATA_CONTROLLER] + programs_coordinator = hass.data[DOMAIN][entry.entry_id][DATA_COORDINATOR][ 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] = [ RainMachineProgram(