jeans/core
jeans/core
1
0
Fork 0
mirror of https://github.com/home-assistant/core.git synced 2025-07-21 12:17:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

Modify Guardian to store a single dataclass in hass.data (#75454)

Browse Source 
* Modify Guardian to store a single dataclass in `hass.data`

* Clarity is better

* Allow entry unload to cancel task
This commit is contained in:
Aaron Bach 2022-07-21 20:32:42 -06:00 committed by GitHub
parent 67e16d77e8
commit b0261dd2eb
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
8 changed files with 112 additions and 149 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

162
homeassistant/components/guardian/__init__.py
View File

@ -35,9 +35,6 @@ from .const import (
API_VALVE_STATUS,
API_WIFI_STATUS,
CONF_UID,
DATA_CLIENT,
DATA_COORDINATOR,
DATA_COORDINATOR_PAIRED_SENSOR,
DOMAIN,
LOGGER,
SIGNAL_PAIRED_SENSOR_COORDINATOR_ADDED,
@ -89,6 +86,16 @@ SERVICE_UPGRADE_FIRMWARE_SCHEMA = vol.Schema(
PLATFORMS = [Platform.BINARY_SENSOR, Platform.BUTTON, Platform.SENSOR, Platform.SWITCH]
@dataclass
class GuardianData:
"""Define an object to be stored in `hass.data`."""
entry: ConfigEntry
client: Client
valve_controller_coordinators: dict[str, GuardianDataUpdateCoordinator]
paired_sensor_manager: PairedSensorManager
@callback
def async_get_entry_id_for_service_call(hass: HomeAssistant, call: ServiceCall) -> str:
"""Get the entry ID related to a service call (by device ID)."""
@ -131,7 +138,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
api_lock = asyncio.Lock()
# Set up GuardianDataUpdateCoordinators for the valve controller:
coordinators: dict[str, GuardianDataUpdateCoordinator] = {}
valve_controller_coordinators: dict[str, GuardianDataUpdateCoordinator] = {}
init_valve_controller_tasks = []
for api, api_coro in (
(API_SENSOR_PAIR_DUMP, client.sensor.pair_dump),
@ -140,7 +147,9 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
(API_VALVE_STATUS, client.valve.status),
(API_WIFI_STATUS, client.wifi.status),
):
coordinator = coordinators[api] = GuardianDataUpdateCoordinator(
coordinator = valve_controller_coordinators[
api
] = GuardianDataUpdateCoordinator(
hass,
client=client,
api_name=api,
@ -154,45 +163,38 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
# Set up an object to evaluate each batch of paired sensor UIDs and add/remove
# devices as appropriate:
paired_sensor_manager = PairedSensorManager(hass, entry, client, api_lock)
await paired_sensor_manager.async_process_latest_paired_sensor_uids()
paired_sensor_manager = PairedSensorManager(
hass,
entry,
client,
api_lock,
valve_controller_coordinators[API_SENSOR_PAIR_DUMP],
)
await paired_sensor_manager.async_initialize()
hass.data.setdefault(DOMAIN, {})
hass.data[DOMAIN][entry.entry_id] = {
DATA_CLIENT: client,
DATA_COORDINATOR: coordinators,
DATA_COORDINATOR_PAIRED_SENSOR: {},
DATA_PAIRED_SENSOR_MANAGER: paired_sensor_manager,
}
@callback
def async_process_paired_sensor_uids() -> None:
"""Define a callback for when new paired sensor data is received."""
hass.async_create_task(
paired_sensor_manager.async_process_latest_paired_sensor_uids()
)
coordinators[API_SENSOR_PAIR_DUMP].async_add_listener(
async_process_paired_sensor_uids
hass.data[DOMAIN][entry.entry_id] = GuardianData(
entry=entry,
client=client,
valve_controller_coordinators=valve_controller_coordinators,
paired_sensor_manager=paired_sensor_manager,
)
# Set up all of the Guardian entity platforms:
await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
@callback
def hydrate_with_entry_and_client(func: Callable) -> Callable:
"""Define a decorator to hydrate a method with args based on service call."""
def call_with_data(func: Callable) -> Callable:
"""Hydrate a service call with the appropriate GuardianData object."""
async def wrapper(call: ServiceCall) -> None:
"""Wrap the service function."""
entry_id = async_get_entry_id_for_service_call(hass, call)
client = hass.data[DOMAIN][entry_id][DATA_CLIENT]
entry = hass.config_entries.async_get_entry(entry_id)
assert entry
data = hass.data[DOMAIN][entry_id]
try:
async with client:
await func(call, entry, client)
async with data.client:
await func(call, data)
except GuardianError as err:
raise HomeAssistantError(
f"Error while executing {func.__name__}: {err}"
@ -200,78 +202,58 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
return wrapper
@hydrate_with_entry_and_client
async def async_disable_ap(
call: ServiceCall, entry: ConfigEntry, client: Client
) -> None:
@call_with_data
async def async_disable_ap(call: ServiceCall, data: GuardianData) -> None:
"""Disable the onboard AP."""
await client.wifi.disable_ap()
await data.client.wifi.disable_ap()
@hydrate_with_entry_and_client
async def async_enable_ap(
call: ServiceCall, entry: ConfigEntry, client: Client
) -> None:
@call_with_data
async def async_enable_ap(call: ServiceCall, data: GuardianData) -> None:
"""Enable the onboard AP."""
await client.wifi.enable_ap()
await data.client.wifi.enable_ap()
@hydrate_with_entry_and_client
async def async_pair_sensor(
call: ServiceCall, entry: ConfigEntry, client: Client
) -> None:
@call_with_data
async def async_pair_sensor(call: ServiceCall, data: GuardianData) -> None:
"""Add a new paired sensor."""
paired_sensor_manager = hass.data[DOMAIN][entry.entry_id][
DATA_PAIRED_SENSOR_MANAGER
]
uid = call.data[CONF_UID]
await data.client.sensor.pair_sensor(uid)
await data.paired_sensor_manager.async_pair_sensor(uid)
await client.sensor.pair_sensor(uid)
await paired_sensor_manager.async_pair_sensor(uid)
@hydrate_with_entry_and_client
async def async_reboot(
call: ServiceCall, entry: ConfigEntry, client: Client
) -> None:
@call_with_data
async def async_reboot(call: ServiceCall, data: GuardianData) -> None:
"""Reboot the valve controller."""
async_log_deprecated_service_call(
hass,
call,
"button.press",
f"button.guardian_valve_controller_{entry.data[CONF_UID]}_reboot",
f"button.guardian_valve_controller_{data.entry.data[CONF_UID]}_reboot",
)
await client.system.reboot()
await data.client.system.reboot()
@hydrate_with_entry_and_client
@call_with_data
async def async_reset_valve_diagnostics(
call: ServiceCall, entry: ConfigEntry, client: Client
call: ServiceCall, data: GuardianData
) -> None:
"""Fully reset system motor diagnostics."""
async_log_deprecated_service_call(
hass,
call,
"button.press",
f"button.guardian_valve_controller_{entry.data[CONF_UID]}_reset_valve_diagnostics",
f"button.guardian_valve_controller_{data.entry.data[CONF_UID]}_reset_valve_diagnostics",
)
await client.valve.reset()
await data.client.valve.reset()
@hydrate_with_entry_and_client
async def async_unpair_sensor(
call: ServiceCall, entry: ConfigEntry, client: Client
) -> None:
@call_with_data
async def async_unpair_sensor(call: ServiceCall, data: GuardianData) -> None:
"""Remove a paired sensor."""
paired_sensor_manager = hass.data[DOMAIN][entry.entry_id][
DATA_PAIRED_SENSOR_MANAGER
]
uid = call.data[CONF_UID]
await data.client.sensor.unpair_sensor(uid)
await data.paired_sensor_manager.async_unpair_sensor(uid)
await client.sensor.unpair_sensor(uid)
await paired_sensor_manager.async_unpair_sensor(uid)
@hydrate_with_entry_and_client
async def async_upgrade_firmware(
call: ServiceCall, entry: ConfigEntry, client: Client
) -> None:
@call_with_data
async def async_upgrade_firmware(call: ServiceCall, data: GuardianData) -> None:
"""Upgrade the device firmware."""
await client.system.upgrade_firmware(
await data.client.system.upgrade_firmware(
url=call.data[CONF_URL],
port=call.data[CONF_PORT],
filename=call.data[CONF_FILENAME],
@ -338,6 +320,7 @@ class PairedSensorManager:
entry: ConfigEntry,
client: Client,
api_lock: asyncio.Lock,
sensor_pair_dump_coordinator: GuardianDataUpdateCoordinator,
) -> None:
"""Initialize."""
self._api_lock = api_lock
@ -345,6 +328,21 @@ class PairedSensorManager:
self._entry = entry
self._hass = hass
self._paired_uids: set[str] = set()
self._sensor_pair_dump_coordinator = sensor_pair_dump_coordinator
self.coordinators: dict[str, GuardianDataUpdateCoordinator] = {}
async def async_initialize(self) -> None:
"""Initialize the manager."""
@callback
def async_create_process_task() -> None:
"""Define a callback for when new paired sensor data is received."""
self._hass.async_create_task(self.async_process_latest_paired_sensor_uids())
cancel_process_task = self._sensor_pair_dump_coordinator.async_add_listener(
async_create_process_task
)
self._entry.async_on_unload(cancel_process_task)
async def async_pair_sensor(self, uid: str) -> None:
"""Add a new paired sensor coordinator."""
@ -352,9 +350,7 @@ class PairedSensorManager:
self._paired_uids.add(uid)
coordinator = self._hass.data[DOMAIN][self._entry.entry_id][
DATA_COORDINATOR_PAIRED_SENSOR
][uid] = GuardianDataUpdateCoordinator(
coordinator = self.coordinators[uid] = GuardianDataUpdateCoordinator(
self._hass,
client=self._client,
api_name=f"{API_SENSOR_PAIRED_SENSOR_STATUS}_{uid}",
@ -375,11 +371,7 @@ class PairedSensorManager:
async def async_process_latest_paired_sensor_uids(self) -> None:
"""Process a list of new UIDs."""
try:
uids = set(
self._hass.data[DOMAIN][self._entry.entry_id][DATA_COORDINATOR][
API_SENSOR_PAIR_DUMP
].data["paired_uids"]
)
uids = set(self._sensor_pair_dump_coordinator.data["paired_uids"])
except KeyError:
# Sometimes the paired_uids key can fail to exist; the user can't do anything
# about it, so in this case, we quietly abort and return:
@ -403,9 +395,7 @@ class PairedSensorManager:
# Clear out objects related to this paired sensor:
self._paired_uids.remove(uid)
self._hass.data[DOMAIN][self._entry.entry_id][
DATA_COORDINATOR_PAIRED_SENSOR
].pop(uid)
self.coordinators.pop(uid)
# Remove the paired sensor device from the device registry (which will
# clean up entities and the entity registry):

15
homeassistant/components/guardian/binary_sensor.py
View File

@ -15,6 +15,7 @@ from homeassistant.helpers.entity import EntityCategory
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from . import (
GuardianData,
PairedSensorEntity,
ValveControllerEntity,
ValveControllerEntityDescription,
@ -23,8 +24,6 @@ from .const import (
API_SYSTEM_ONBOARD_SENSOR_STATUS,
API_WIFI_STATUS,
CONF_UID,
DATA_COORDINATOR,
DATA_COORDINATOR_PAIRED_SENSOR,
DOMAIN,
SIGNAL_PAIRED_SENSOR_COORDINATOR_ADDED,
)
@ -79,16 +78,14 @@ async def async_setup_entry(
hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
) -> None:
"""Set up Guardian switches based on a config entry."""
entry_data = hass.data[DOMAIN][entry.entry_id]
paired_sensor_coordinators = entry_data[DATA_COORDINATOR_PAIRED_SENSOR]
valve_controller_coordinators = entry_data[DATA_COORDINATOR]
data: GuardianData = hass.data[DOMAIN][entry.entry_id]
@callback
def add_new_paired_sensor(uid: str) -> None:
"""Add a new paired sensor."""
async_add_entities(
PairedSensorBinarySensor(
entry, paired_sensor_coordinators[uid], description
entry, data.paired_sensor_manager.coordinators[uid], description
)
for description in PAIRED_SENSOR_DESCRIPTIONS
)
@ -104,7 +101,9 @@ async def async_setup_entry(
# Add all valve controller-specific binary sensors:
sensors: list[PairedSensorBinarySensor | ValveControllerBinarySensor] = [
ValveControllerBinarySensor(entry, valve_controller_coordinators, description)
ValveControllerBinarySensor(
entry, data.valve_controller_coordinators, description
)
for description in VALVE_CONTROLLER_DESCRIPTIONS
]
@ -112,7 +111,7 @@ async def async_setup_entry(
sensors.extend(
[
PairedSensorBinarySensor(entry, coordinator, description)
for coordinator in paired_sensor_coordinators.values()
for coordinator in data.paired_sensor_manager.coordinators.values()
for description in PAIRED_SENSOR_DESCRIPTIONS
]
)

19
homeassistant/components/guardian/button.py
View File

@ -18,9 +18,8 @@ from homeassistant.exceptions import HomeAssistantError
from homeassistant.helpers.entity import EntityCategory
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from . import ValveControllerEntity, ValveControllerEntityDescription
from .const import API_SYSTEM_DIAGNOSTICS, DATA_CLIENT, DATA_COORDINATOR, DOMAIN
from .util import GuardianDataUpdateCoordinator
from . import GuardianData, ValveControllerEntity, ValveControllerEntityDescription
from .const import API_SYSTEM_DIAGNOSTICS, DOMAIN
@dataclass
@ -77,13 +76,10 @@ async def async_setup_entry(
hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
) -> None:
"""Set up Guardian buttons based on a config entry."""
entry_data = hass.data[DOMAIN][entry.entry_id]
client = entry_data[DATA_CLIENT]
valve_controller_coordinators = entry_data[DATA_COORDINATOR]
data: GuardianData = hass.data[DOMAIN][entry.entry_id]
async_add_entities(
GuardianButton(entry, valve_controller_coordinators, description, client)
for description in BUTTON_DESCRIPTIONS
GuardianButton(entry, data, description) for description in BUTTON_DESCRIPTIONS
)
@ -98,14 +94,13 @@ class GuardianButton(ValveControllerEntity, ButtonEntity):
def __init__(
self,
entry: ConfigEntry,
coordinators: dict[str, GuardianDataUpdateCoordinator],
data: GuardianData,
description: ValveControllerButtonDescription,
client: Client,
) -> None:
"""Initialize."""
super().__init__(entry, coordinators, description)
super().__init__(entry, data.valve_controller_coordinators, description)
self._client = client
self._client = data.client
async def async_press(self) -> None:
"""Send out a restart command."""

4
homeassistant/components/guardian/const.py
View File

@ -14,8 +14,4 @@ API_WIFI_STATUS = "wifi_status"
CONF_UID = "uid"
DATA_CLIENT = "client"
DATA_COORDINATOR = "coordinator"
DATA_COORDINATOR_PAIRED_SENSOR = "coordinator_paired_sensor"
SIGNAL_PAIRED_SENSOR_COORDINATOR_ADDED = "guardian_paired_sensor_coordinator_added_{0}"

15
homeassistant/components/guardian/diagnostics.py
View File

@ -7,8 +7,8 @@ from homeassistant.components.diagnostics import async_redact_data
from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant
from .const import CONF_UID, DATA_COORDINATOR, DATA_COORDINATOR_PAIRED_SENSOR, DOMAIN
from .util import GuardianDataUpdateCoordinator
from . import GuardianData
from .const import CONF_UID, DOMAIN
CONF_BSSID = "bssid"
CONF_PAIRED_UIDS = "paired_uids"
@ -26,12 +26,7 @@ async def async_get_config_entry_diagnostics(
hass: HomeAssistant, entry: ConfigEntry
) -> dict[str, Any]:
"""Return diagnostics for a config entry."""
data = hass.data[DOMAIN][entry.entry_id]
coordinators: dict[str, GuardianDataUpdateCoordinator] = data[DATA_COORDINATOR]
paired_sensor_coordinators: dict[str, GuardianDataUpdateCoordinator] = data[
DATA_COORDINATOR_PAIRED_SENSOR
]
data: GuardianData = hass.data[DOMAIN][entry.entry_id]
return {
"entry": {
@ -41,11 +36,11 @@ async def async_get_config_entry_diagnostics(
"data": {
"valve_controller": {
api_category: async_redact_data(coordinator.data, TO_REDACT)
for api_category, coordinator in coordinators.items()
for api_category, coordinator in data.valve_controller_coordinators.items()
},
"paired_sensors": [
async_redact_data(coordinator.data, TO_REDACT)
for coordinator in paired_sensor_coordinators.values()
for coordinator in data.paired_sensor_manager.coordinators.values()
],
},
}

15
homeassistant/components/guardian/sensor.py
View File

@ -17,6 +17,7 @@ from homeassistant.helpers.entity import EntityCategory
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from . import (
GuardianData,
PairedSensorEntity,
ValveControllerEntity,
ValveControllerEntityDescription,
@ -25,8 +26,6 @@ from .const import (
API_SYSTEM_DIAGNOSTICS,
API_SYSTEM_ONBOARD_SENSOR_STATUS,
CONF_UID,
DATA_COORDINATOR,
DATA_COORDINATOR_PAIRED_SENSOR,
DOMAIN,
SIGNAL_PAIRED_SENSOR_COORDINATOR_ADDED,
)
@ -83,15 +82,15 @@ async def async_setup_entry(
hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
) -> None:
"""Set up Guardian switches based on a config entry."""
entry_data = hass.data[DOMAIN][entry.entry_id]
paired_sensor_coordinators = entry_data[DATA_COORDINATOR_PAIRED_SENSOR]
valve_controller_coordinators = entry_data[DATA_COORDINATOR]
data: GuardianData = hass.data[DOMAIN][entry.entry_id]
@callback
def add_new_paired_sensor(uid: str) -> None:
"""Add a new paired sensor."""
async_add_entities(
PairedSensorSensor(entry, paired_sensor_coordinators[uid], description)
PairedSensorSensor(
entry, data.paired_sensor_manager.coordinators[uid], description
)
for description in PAIRED_SENSOR_DESCRIPTIONS
)
@ -106,7 +105,7 @@ async def async_setup_entry(
# Add all valve controller-specific binary sensors:
sensors: list[PairedSensorSensor | ValveControllerSensor] = [
ValveControllerSensor(entry, valve_controller_coordinators, description)
ValveControllerSensor(entry, data.valve_controller_coordinators, description)
for description in VALVE_CONTROLLER_DESCRIPTIONS
]
@ -114,7 +113,7 @@ async def async_setup_entry(
sensors.extend(
[
PairedSensorSensor(entry, coordinator, description)
for coordinator in paired_sensor_coordinators.values()
for coordinator in data.paired_sensor_manager.coordinators.values()
for description in PAIRED_SENSOR_DESCRIPTIONS
]
)

19
homeassistant/components/guardian/switch.py
View File

@ -4,7 +4,6 @@ from __future__ import annotations
from dataclasses import dataclass
from typing import Any
from aioguardian import Client
from aioguardian.errors import GuardianError
from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription
@ -13,9 +12,8 @@ from homeassistant.core import HomeAssistant, callback
from homeassistant.exceptions import HomeAssistantError
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from . import ValveControllerEntity, ValveControllerEntityDescription
from .const import API_VALVE_STATUS, DATA_CLIENT, DATA_COORDINATOR, DOMAIN
from .util import GuardianDataUpdateCoordinator
from . import GuardianData, ValveControllerEntity, ValveControllerEntityDescription
from .const import API_VALVE_STATUS, DOMAIN
ATTR_AVG_CURRENT = "average_current"
ATTR_INST_CURRENT = "instantaneous_current"
@ -46,12 +44,10 @@ async def async_setup_entry(
hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
) -> None:
"""Set up Guardian switches based on a config entry."""
entry_data = hass.data[DOMAIN][entry.entry_id]
client = entry_data[DATA_CLIENT]
valve_controller_coordinators = entry_data[DATA_COORDINATOR]
data: GuardianData = hass.data[DOMAIN][entry.entry_id]
async_add_entities(
ValveControllerSwitch(entry, valve_controller_coordinators, description, client)
ValveControllerSwitch(entry, data, description)
for description in VALVE_CONTROLLER_DESCRIPTIONS
)
@ -71,15 +67,14 @@ class ValveControllerSwitch(ValveControllerEntity, SwitchEntity):
def __init__(
self,
entry: ConfigEntry,
coordinators: dict[str, GuardianDataUpdateCoordinator],
data: GuardianData,
description: ValveControllerSwitchDescription,
client: Client,
) -> None:
"""Initialize."""
super().__init__(entry, coordinators, description)
super().__init__(entry, data.valve_controller_coordinators, description)
self._attr_is_on = True
self._client = client
self._client = data.client
@callback
def _async_update_from_latest_data(self) -> None:

12
tests/components/guardian/test_diagnostics.py
View File

@ -1,22 +1,16 @@
"""Test Guardian diagnostics."""
from homeassistant.components.diagnostics import REDACTED
from homeassistant.components.guardian import (
DATA_PAIRED_SENSOR_MANAGER,
DOMAIN,
PairedSensorManager,
)
from homeassistant.components.guardian import DOMAIN, GuardianData
from tests.components.diagnostics import get_diagnostics_for_config_entry
async def test_entry_diagnostics(hass, config_entry, hass_client, setup_guardian):
"""Test config entry diagnostics."""
paired_sensor_manager: PairedSensorManager = hass.data[DOMAIN][
config_entry.entry_id
][DATA_PAIRED_SENSOR_MANAGER]
data: GuardianData = hass.data[DOMAIN][config_entry.entry_id]
# Simulate the pairing of a paired sensor:
await paired_sensor_manager.async_pair_sensor("AABBCCDDEEFF")
await data.paired_sensor_manager.async_pair_sensor("AABBCCDDEEFF")
assert await get_diagnostics_for_config_entry(hass, hass_client, config_entry) == {
"entry": {