diff --git a/homeassistant/components/guardian/__init__.py b/homeassistant/components/guardian/__init__.py
index d15b7d57aea9eb1c0307e21506090d614e330d5b..da22066d7aa98f4caaba38d1b0eb36dac154b6bd 100644
--- a/homeassistant/components/guardian/__init__.py
+++ b/homeassistant/components/guardian/__init__.py
@@ -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):
diff --git a/homeassistant/components/guardian/binary_sensor.py b/homeassistant/components/guardian/binary_sensor.py
index 9b824ab589fb578369756a1ab5527472d28fe869..eb6d49c3ec133f79561198459e174cc5ee6e9e1c 100644
--- a/homeassistant/components/guardian/binary_sensor.py
+++ b/homeassistant/components/guardian/binary_sensor.py
@@ -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
]
)
diff --git a/homeassistant/components/guardian/button.py b/homeassistant/components/guardian/button.py
index e013cde85d6fe3111988a72af09d44a851ef02d7..01efb7deba4e415a7f41155b7d4df75d2c757733 100644
--- a/homeassistant/components/guardian/button.py
+++ b/homeassistant/components/guardian/button.py
@@ -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."""
diff --git a/homeassistant/components/guardian/const.py b/homeassistant/components/guardian/const.py
index 3499db24c037b21c3a07ca1f6702194967f56a3a..c7d025ba712173caedf6aad5180a038a9c391a87 100644
--- a/homeassistant/components/guardian/const.py
+++ b/homeassistant/components/guardian/const.py
@@ -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}"
diff --git a/homeassistant/components/guardian/diagnostics.py b/homeassistant/components/guardian/diagnostics.py
index 175136b33f4657dc2bb3df45bd437462ffd4e20a..d53dcb68fa8cb53690740d4af27f7da565b290c4 100644
--- a/homeassistant/components/guardian/diagnostics.py
+++ b/homeassistant/components/guardian/diagnostics.py
@@ -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()
],
},
}
diff --git a/homeassistant/components/guardian/sensor.py b/homeassistant/components/guardian/sensor.py
index c4fd1e110fa06773e6e67dd88369d3693e4b40ca..bf7d9e7122ae13306cda3b261e6c6c1cdd157df3 100644
--- a/homeassistant/components/guardian/sensor.py
+++ b/homeassistant/components/guardian/sensor.py
@@ -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
]
)
diff --git a/homeassistant/components/guardian/switch.py b/homeassistant/components/guardian/switch.py
index c58f4548a87dcd2e9adfe033a56439b533f62d34..4e100ce4fe4ebbd887461515f44c6b305611176c 100644
--- a/homeassistant/components/guardian/switch.py
+++ b/homeassistant/components/guardian/switch.py
@@ -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:
diff --git a/tests/components/guardian/test_diagnostics.py b/tests/components/guardian/test_diagnostics.py
index f48c988c907fc12b8b4764d83b685cbdd9e3c867..2269d09b1eba96e6d4e090aa1a87ebe942191faa 100644
--- a/tests/components/guardian/test_diagnostics.py
+++ b/tests/components/guardian/test_diagnostics.py
@@ -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": {