jeans/core
jeans/core
1
0
Fork 0
mirror of https://github.com/home-assistant/core.git synced 2025-07-23 13:17:32 +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_VALVE_STATUS,
API_WIFI_STATUS, API_WIFI_STATUS,
CONF_UID, CONF_UID,
DATA_CLIENT,
DATA_COORDINATOR,
DATA_COORDINATOR_PAIRED_SENSOR,
DOMAIN, DOMAIN,
LOGGER, LOGGER,
SIGNAL_PAIRED_SENSOR_COORDINATOR_ADDED, 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] 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 @callback
def async_get_entry_id_for_service_call(hass: HomeAssistant, call: ServiceCall) -> str: 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).""" """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() api_lock = asyncio.Lock()
# Set up GuardianDataUpdateCoordinators for the valve controller: # Set up GuardianDataUpdateCoordinators for the valve controller:
coordinators: dict[str, GuardianDataUpdateCoordinator] = {} valve_controller_coordinators: dict[str, GuardianDataUpdateCoordinator] = {}
init_valve_controller_tasks = [] init_valve_controller_tasks = []
for api, api_coro in ( for api, api_coro in (
(API_SENSOR_PAIR_DUMP, client.sensor.pair_dump), (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_VALVE_STATUS, client.valve.status),
(API_WIFI_STATUS, client.wifi.status), (API_WIFI_STATUS, client.wifi.status),
): ):
coordinator = coordinators[api] = GuardianDataUpdateCoordinator( coordinator = valve_controller_coordinators[
api
] = GuardianDataUpdateCoordinator(
hass, hass,
client=client, client=client,
api_name=api, 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 # Set up an object to evaluate each batch of paired sensor UIDs and add/remove
# devices as appropriate: # devices as appropriate:
paired_sensor_manager = PairedSensorManager(hass, entry, client, api_lock) paired_sensor_manager = PairedSensorManager(
await paired_sensor_manager.async_process_latest_paired_sensor_uids() hass,
entry,
client,
api_lock,
valve_controller_coordinators[API_SENSOR_PAIR_DUMP],
)
await paired_sensor_manager.async_initialize()
hass.data.setdefault(DOMAIN, {}) hass.data.setdefault(DOMAIN, {})
hass.data[DOMAIN][entry.entry_id] = { hass.data[DOMAIN][entry.entry_id] = GuardianData(
DATA_CLIENT: client, entry=entry,
DATA_COORDINATOR: coordinators, client=client,
DATA_COORDINATOR_PAIRED_SENSOR: {}, valve_controller_coordinators=valve_controller_coordinators,
DATA_PAIRED_SENSOR_MANAGER: paired_sensor_manager, 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
) )
# Set up all of the Guardian entity platforms: # Set up all of the Guardian entity platforms:
await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
@callback @callback
def hydrate_with_entry_and_client(func: Callable) -> Callable: def call_with_data(func: Callable) -> Callable:
"""Define a decorator to hydrate a method with args based on service call.""" """Hydrate a service call with the appropriate GuardianData object."""
async def wrapper(call: ServiceCall) -> None: async def wrapper(call: ServiceCall) -> None:
"""Wrap the service function.""" """Wrap the service function."""
entry_id = async_get_entry_id_for_service_call(hass, call) entry_id = async_get_entry_id_for_service_call(hass, call)
client = hass.data[DOMAIN][entry_id][DATA_CLIENT] data = hass.data[DOMAIN][entry_id]
entry = hass.config_entries.async_get_entry(entry_id)
assert entry
try: try:
async with client: async with data.client:
await func(call, entry, client) await func(call, data)
except GuardianError as err: except GuardianError as err:
raise HomeAssistantError( raise HomeAssistantError(
f"Error while executing {func.__name__}: {err}" f"Error while executing {func.__name__}: {err}"
@ -200,78 +202,58 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
return wrapper return wrapper
@hydrate_with_entry_and_client @call_with_data
async def async_disable_ap( async def async_disable_ap(call: ServiceCall, data: GuardianData) -> None:
call: ServiceCall, entry: ConfigEntry, client: Client
) -> None:
"""Disable the onboard AP.""" """Disable the onboard AP."""
await client.wifi.disable_ap() await data.client.wifi.disable_ap()
@hydrate_with_entry_and_client @call_with_data
async def async_enable_ap( async def async_enable_ap(call: ServiceCall, data: GuardianData) -> None:
call: ServiceCall, entry: ConfigEntry, client: Client
) -> None:
"""Enable the onboard AP.""" """Enable the onboard AP."""
await client.wifi.enable_ap() await data.client.wifi.enable_ap()
@hydrate_with_entry_and_client @call_with_data
async def async_pair_sensor( async def async_pair_sensor(call: ServiceCall, data: GuardianData) -> None:
call: ServiceCall, entry: ConfigEntry, client: Client
) -> None:
"""Add a new paired sensor.""" """Add a new paired sensor."""
paired_sensor_manager = hass.data[DOMAIN][entry.entry_id][
DATA_PAIRED_SENSOR_MANAGER
]
uid = call.data[CONF_UID] 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) @call_with_data
await paired_sensor_manager.async_pair_sensor(uid) async def async_reboot(call: ServiceCall, data: GuardianData) -> None:
@hydrate_with_entry_and_client
async def async_reboot(
call: ServiceCall, entry: ConfigEntry, client: Client
) -> None:
"""Reboot the valve controller.""" """Reboot the valve controller."""
async_log_deprecated_service_call( async_log_deprecated_service_call(
hass, hass,
call, call,
"button.press", "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( async def async_reset_valve_diagnostics(
call: ServiceCall, entry: ConfigEntry, client: Client call: ServiceCall, data: GuardianData
) -> None: ) -> None:
"""Fully reset system motor diagnostics.""" """Fully reset system motor diagnostics."""
async_log_deprecated_service_call( async_log_deprecated_service_call(
hass, hass,
call, call,
"button.press", "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 @call_with_data
async def async_unpair_sensor( async def async_unpair_sensor(call: ServiceCall, data: GuardianData) -> None:
call: ServiceCall, entry: ConfigEntry, client: Client
) -> None:
"""Remove a paired sensor.""" """Remove a paired sensor."""
paired_sensor_manager = hass.data[DOMAIN][entry.entry_id][
DATA_PAIRED_SENSOR_MANAGER
]
uid = call.data[CONF_UID] 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) @call_with_data
await paired_sensor_manager.async_unpair_sensor(uid) async def async_upgrade_firmware(call: ServiceCall, data: GuardianData) -> None:
@hydrate_with_entry_and_client
async def async_upgrade_firmware(
call: ServiceCall, entry: ConfigEntry, client: Client
) -> None:
"""Upgrade the device firmware.""" """Upgrade the device firmware."""
await client.system.upgrade_firmware( await data.client.system.upgrade_firmware(
url=call.data[CONF_URL], url=call.data[CONF_URL],
port=call.data[CONF_PORT], port=call.data[CONF_PORT],
filename=call.data[CONF_FILENAME], filename=call.data[CONF_FILENAME],
@ -338,6 +320,7 @@ class PairedSensorManager:
entry: ConfigEntry, entry: ConfigEntry,
client: Client, client: Client,
api_lock: asyncio.Lock, api_lock: asyncio.Lock,
sensor_pair_dump_coordinator: GuardianDataUpdateCoordinator,
) -> None: ) -> None:
"""Initialize.""" """Initialize."""
self._api_lock = api_lock self._api_lock = api_lock
@ -345,6 +328,21 @@ class PairedSensorManager:
self._entry = entry self._entry = entry
self._hass = hass self._hass = hass
self._paired_uids: set[str] = set() 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: async def async_pair_sensor(self, uid: str) -> None:
"""Add a new paired sensor coordinator.""" """Add a new paired sensor coordinator."""
@ -352,9 +350,7 @@ class PairedSensorManager:
self._paired_uids.add(uid) self._paired_uids.add(uid)
coordinator = self._hass.data[DOMAIN][self._entry.entry_id][ coordinator = self.coordinators[uid] = GuardianDataUpdateCoordinator(
DATA_COORDINATOR_PAIRED_SENSOR
][uid] = GuardianDataUpdateCoordinator(
self._hass, self._hass,
client=self._client, client=self._client,
api_name=f"{API_SENSOR_PAIRED_SENSOR_STATUS}_{uid}", 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: async def async_process_latest_paired_sensor_uids(self) -> None:
"""Process a list of new UIDs.""" """Process a list of new UIDs."""
try: try:
uids = set( uids = set(self._sensor_pair_dump_coordinator.data["paired_uids"])
self._hass.data[DOMAIN][self._entry.entry_id][DATA_COORDINATOR][
API_SENSOR_PAIR_DUMP
].data["paired_uids"]
)
except KeyError: except KeyError:
# Sometimes the paired_uids key can fail to exist; the user can't do anything # 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: # 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: # Clear out objects related to this paired sensor:
self._paired_uids.remove(uid) self._paired_uids.remove(uid)
self._hass.data[DOMAIN][self._entry.entry_id][ self.coordinators.pop(uid)
DATA_COORDINATOR_PAIRED_SENSOR
].pop(uid)
# Remove the paired sensor device from the device registry (which will # Remove the paired sensor device from the device registry (which will
# clean up entities and the entity registry): # 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 homeassistant.helpers.entity_platform import AddEntitiesCallback
from . import ( from . import (
GuardianData,
PairedSensorEntity, PairedSensorEntity,
ValveControllerEntity, ValveControllerEntity,
ValveControllerEntityDescription, ValveControllerEntityDescription,
@ -23,8 +24,6 @@ from .const import (
API_SYSTEM_ONBOARD_SENSOR_STATUS, API_SYSTEM_ONBOARD_SENSOR_STATUS,
API_WIFI_STATUS, API_WIFI_STATUS,
CONF_UID, CONF_UID,
DATA_COORDINATOR,
DATA_COORDINATOR_PAIRED_SENSOR,
DOMAIN, DOMAIN,
SIGNAL_PAIRED_SENSOR_COORDINATOR_ADDED, SIGNAL_PAIRED_SENSOR_COORDINATOR_ADDED,
) )
@ -79,16 +78,14 @@ 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 Guardian switches based on a config entry.""" """Set up Guardian switches based on a config entry."""
entry_data = hass.data[DOMAIN][entry.entry_id] data: GuardianData = hass.data[DOMAIN][entry.entry_id]
paired_sensor_coordinators = entry_data[DATA_COORDINATOR_PAIRED_SENSOR]
valve_controller_coordinators = entry_data[DATA_COORDINATOR]
@callback @callback
def add_new_paired_sensor(uid: str) -> None: def add_new_paired_sensor(uid: str) -> None:
"""Add a new paired sensor.""" """Add a new paired sensor."""
async_add_entities( async_add_entities(
PairedSensorBinarySensor( PairedSensorBinarySensor(
entry, paired_sensor_coordinators[uid], description entry, data.paired_sensor_manager.coordinators[uid], description
) )
for description in PAIRED_SENSOR_DESCRIPTIONS for description in PAIRED_SENSOR_DESCRIPTIONS
) )
@ -104,7 +101,9 @@ async def async_setup_entry(
# Add all valve controller-specific binary sensors: # Add all valve controller-specific binary sensors:
sensors: list[PairedSensorBinarySensor | ValveControllerBinarySensor] = [ sensors: list[PairedSensorBinarySensor | ValveControllerBinarySensor] = [
ValveControllerBinarySensor(entry, valve_controller_coordinators, description) ValveControllerBinarySensor(
entry, data.valve_controller_coordinators, description
)
for description in VALVE_CONTROLLER_DESCRIPTIONS for description in VALVE_CONTROLLER_DESCRIPTIONS
] ]
@ -112,7 +111,7 @@ async def async_setup_entry(
sensors.extend( sensors.extend(
[ [
PairedSensorBinarySensor(entry, coordinator, description) 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 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 import EntityCategory
from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.entity_platform import AddEntitiesCallback
from . import ValveControllerEntity, ValveControllerEntityDescription from . import GuardianData, ValveControllerEntity, ValveControllerEntityDescription
from .const import API_SYSTEM_DIAGNOSTICS, DATA_CLIENT, DATA_COORDINATOR, DOMAIN from .const import API_SYSTEM_DIAGNOSTICS, DOMAIN
from .util import GuardianDataUpdateCoordinator
@dataclass @dataclass
@ -77,13 +76,10 @@ 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 Guardian buttons based on a config entry.""" """Set up Guardian buttons based on a config entry."""
entry_data = hass.data[DOMAIN][entry.entry_id] data: GuardianData = hass.data[DOMAIN][entry.entry_id]
client = entry_data[DATA_CLIENT]
valve_controller_coordinators = entry_data[DATA_COORDINATOR]
async_add_entities( async_add_entities(
GuardianButton(entry, valve_controller_coordinators, description, client) GuardianButton(entry, data, description) for description in BUTTON_DESCRIPTIONS
for description in BUTTON_DESCRIPTIONS
) )
@ -98,14 +94,13 @@ class GuardianButton(ValveControllerEntity, ButtonEntity):
def __init__( def __init__(
self, self,
entry: ConfigEntry, entry: ConfigEntry,
coordinators: dict[str, GuardianDataUpdateCoordinator], data: GuardianData,
description: ValveControllerButtonDescription, description: ValveControllerButtonDescription,
client: Client,
) -> None: ) -> None:
"""Initialize.""" """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: async def async_press(self) -> None:
"""Send out a restart command.""" """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" 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}" 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.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant from homeassistant.core import HomeAssistant
from .const import CONF_UID, DATA_COORDINATOR, DATA_COORDINATOR_PAIRED_SENSOR, DOMAIN from . import GuardianData
from .util import GuardianDataUpdateCoordinator from .const import CONF_UID, DOMAIN
CONF_BSSID = "bssid" CONF_BSSID = "bssid"
CONF_PAIRED_UIDS = "paired_uids" CONF_PAIRED_UIDS = "paired_uids"
@ -26,12 +26,7 @@ async def async_get_config_entry_diagnostics(
hass: HomeAssistant, entry: ConfigEntry hass: HomeAssistant, entry: ConfigEntry
) -> dict[str, Any]: ) -> dict[str, Any]:
"""Return diagnostics for a config entry.""" """Return diagnostics for a config entry."""
data = hass.data[DOMAIN][entry.entry_id] data: GuardianData = hass.data[DOMAIN][entry.entry_id]
coordinators: dict[str, GuardianDataUpdateCoordinator] = data[DATA_COORDINATOR]
paired_sensor_coordinators: dict[str, GuardianDataUpdateCoordinator] = data[
DATA_COORDINATOR_PAIRED_SENSOR
]
return { return {
"entry": { "entry": {
@ -41,11 +36,11 @@ async def async_get_config_entry_diagnostics(
"data": { "data": {
"valve_controller": { "valve_controller": {
api_category: async_redact_data(coordinator.data, TO_REDACT) 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": [ "paired_sensors": [
async_redact_data(coordinator.data, TO_REDACT) 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 homeassistant.helpers.entity_platform import AddEntitiesCallback
from . import ( from . import (
GuardianData,
PairedSensorEntity, PairedSensorEntity,
ValveControllerEntity, ValveControllerEntity,
ValveControllerEntityDescription, ValveControllerEntityDescription,
@ -25,8 +26,6 @@ from .const import (
API_SYSTEM_DIAGNOSTICS, API_SYSTEM_DIAGNOSTICS,
API_SYSTEM_ONBOARD_SENSOR_STATUS, API_SYSTEM_ONBOARD_SENSOR_STATUS,
CONF_UID, CONF_UID,
DATA_COORDINATOR,
DATA_COORDINATOR_PAIRED_SENSOR,
DOMAIN, DOMAIN,
SIGNAL_PAIRED_SENSOR_COORDINATOR_ADDED, SIGNAL_PAIRED_SENSOR_COORDINATOR_ADDED,
) )
@ -83,15 +82,15 @@ 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 Guardian switches based on a config entry.""" """Set up Guardian switches based on a config entry."""
entry_data = hass.data[DOMAIN][entry.entry_id] data: GuardianData = hass.data[DOMAIN][entry.entry_id]
paired_sensor_coordinators = entry_data[DATA_COORDINATOR_PAIRED_SENSOR]
valve_controller_coordinators = entry_data[DATA_COORDINATOR]
@callback @callback
def add_new_paired_sensor(uid: str) -> None: def add_new_paired_sensor(uid: str) -> None:
"""Add a new paired sensor.""" """Add a new paired sensor."""
async_add_entities( 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 for description in PAIRED_SENSOR_DESCRIPTIONS
) )
@ -106,7 +105,7 @@ async def async_setup_entry(
# Add all valve controller-specific binary sensors: # Add all valve controller-specific binary sensors:
sensors: list[PairedSensorSensor | ValveControllerSensor] = [ sensors: list[PairedSensorSensor | ValveControllerSensor] = [
ValveControllerSensor(entry, valve_controller_coordinators, description) ValveControllerSensor(entry, data.valve_controller_coordinators, description)
for description in VALVE_CONTROLLER_DESCRIPTIONS for description in VALVE_CONTROLLER_DESCRIPTIONS
] ]
@ -114,7 +113,7 @@ async def async_setup_entry(
sensors.extend( sensors.extend(
[ [
PairedSensorSensor(entry, coordinator, description) 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 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 dataclasses import dataclass
from typing import Any from typing import Any
from aioguardian import Client
from aioguardian.errors import GuardianError from aioguardian.errors import GuardianError
from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription
@ -13,9 +12,8 @@ from homeassistant.core import HomeAssistant, callback
from homeassistant.exceptions import HomeAssistantError from homeassistant.exceptions import HomeAssistantError
from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.entity_platform import AddEntitiesCallback
from . import ValveControllerEntity, ValveControllerEntityDescription from . import GuardianData, ValveControllerEntity, ValveControllerEntityDescription
from .const import API_VALVE_STATUS, DATA_CLIENT, DATA_COORDINATOR, DOMAIN from .const import API_VALVE_STATUS, DOMAIN
from .util import GuardianDataUpdateCoordinator
ATTR_AVG_CURRENT = "average_current" ATTR_AVG_CURRENT = "average_current"
ATTR_INST_CURRENT = "instantaneous_current" ATTR_INST_CURRENT = "instantaneous_current"
@ -46,12 +44,10 @@ 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 Guardian switches based on a config entry.""" """Set up Guardian switches based on a config entry."""
entry_data = hass.data[DOMAIN][entry.entry_id] data: GuardianData = hass.data[DOMAIN][entry.entry_id]
client = entry_data[DATA_CLIENT]
valve_controller_coordinators = entry_data[DATA_COORDINATOR]
async_add_entities( async_add_entities(
ValveControllerSwitch(entry, valve_controller_coordinators, description, client) ValveControllerSwitch(entry, data, description)
for description in VALVE_CONTROLLER_DESCRIPTIONS for description in VALVE_CONTROLLER_DESCRIPTIONS
) )
@ -71,15 +67,14 @@ class ValveControllerSwitch(ValveControllerEntity, SwitchEntity):
def __init__( def __init__(
self, self,
entry: ConfigEntry, entry: ConfigEntry,
coordinators: dict[str, GuardianDataUpdateCoordinator], data: GuardianData,
description: ValveControllerSwitchDescription, description: ValveControllerSwitchDescription,
client: Client,
) -> None: ) -> None:
"""Initialize.""" """Initialize."""
super().__init__(entry, coordinators, description) super().__init__(entry, data.valve_controller_coordinators, description)
self._attr_is_on = True self._attr_is_on = True
self._client = client self._client = data.client
@callback @callback
def _async_update_from_latest_data(self) -> None: def _async_update_from_latest_data(self) -> None:

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

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