From faf2a447a42b91e005d9a91416fabc25e2757701 Mon Sep 17 00:00:00 2001 From: Robert Hillis Date: Mon, 17 Jun 2024 21:43:45 -0400 Subject: [PATCH] Store runtime data inside the config entry in Sense (#119740) --- homeassistant/components/sense/__init__.py | 42 ++++++++++--------- .../components/sense/binary_sensor.py | 25 ++++------- homeassistant/components/sense/const.py | 4 -- homeassistant/components/sense/sensor.py | 21 ++++------ 4 files changed, 37 insertions(+), 55 deletions(-) diff --git a/homeassistant/components/sense/__init__.py b/homeassistant/components/sense/__init__.py index 88af9fa990b..28408c0cb7d 100644 --- a/homeassistant/components/sense/__init__.py +++ b/homeassistant/components/sense/__init__.py @@ -1,7 +1,9 @@ """Support for monitoring a Sense energy sensor.""" +from dataclasses import dataclass from datetime import timedelta import logging +from typing import Any from sense_energy import ( ASyncSenseable, @@ -25,20 +27,16 @@ from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, Upda from .const import ( ACTIVE_UPDATE_RATE, - DOMAIN, SENSE_CONNECT_EXCEPTIONS, - SENSE_DATA, SENSE_DEVICE_UPDATE, - SENSE_DEVICES_DATA, - SENSE_DISCOVERED_DEVICES_DATA, SENSE_TIMEOUT_EXCEPTIONS, - SENSE_TRENDS_COORDINATOR, SENSE_WEBSOCKET_EXCEPTIONS, ) _LOGGER = logging.getLogger(__name__) PLATFORMS = [Platform.BINARY_SENSOR, Platform.SENSOR] +type SenseConfigEntry = ConfigEntry[SenseData] class SenseDevicesData: @@ -57,7 +55,17 @@ class SenseDevicesData: return self._data_by_device.get(sense_device_id) -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +@dataclass(kw_only=True, slots=True) +class SenseData: + """Sense data type.""" + + data: ASyncSenseable + device_data: SenseDevicesData + trends: DataUpdateCoordinator[None] + discovered: list[dict[str, Any]] + + +async def async_setup_entry(hass: HomeAssistant, entry: SenseConfigEntry) -> bool: """Set up Sense from a config entry.""" entry_data = entry.data @@ -91,7 +99,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: except SENSE_CONNECT_EXCEPTIONS as err: raise ConfigEntryNotReady(str(err)) from err - sense_devices_data = SenseDevicesData() try: sense_discovered_devices = await gateway.get_discovered_device_data() await gateway.update_realtime() @@ -132,12 +139,12 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: "sense.trends-coordinator-refresh", ) - hass.data.setdefault(DOMAIN, {})[entry.entry_id] = { - SENSE_DATA: gateway, - SENSE_DEVICES_DATA: sense_devices_data, - SENSE_TRENDS_COORDINATOR: trends_coordinator, - SENSE_DISCOVERED_DEVICES_DATA: sense_discovered_devices, - } + entry.runtime_data = SenseData( + data=gateway, + device_data=SenseDevicesData(), + trends=trends_coordinator, + discovered=sense_discovered_devices, + ) await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) @@ -152,7 +159,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: data = gateway.get_realtime() if "devices" in data: - sense_devices_data.set_devices_data(data["devices"]) + entry.runtime_data.device_data.set_devices_data(data["devices"]) async_dispatcher_send(hass, f"{SENSE_DEVICE_UPDATE}-{gateway.sense_monitor_id}") remove_update_callback = async_track_time_interval( @@ -173,9 +180,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: return True -async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_unload_entry(hass: HomeAssistant, entry: SenseConfigEntry) -> bool: """Unload a config entry.""" - unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS) - if unload_ok: - hass.data[DOMAIN].pop(entry.entry_id) - return unload_ok + return await hass.config_entries.async_unload_platforms(entry, PLATFORMS) diff --git a/homeassistant/components/sense/binary_sensor.py b/homeassistant/components/sense/binary_sensor.py index 7dde4c029b1..5640dd19961 100644 --- a/homeassistant/components/sense/binary_sensor.py +++ b/homeassistant/components/sense/binary_sensor.py @@ -6,40 +6,29 @@ from homeassistant.components.binary_sensor import ( BinarySensorDeviceClass, BinarySensorEntity, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant, callback from homeassistant.helpers import entity_registry as er from homeassistant.helpers.dispatcher import async_dispatcher_connect from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import ( - ATTRIBUTION, - DOMAIN, - MDI_ICONS, - SENSE_DATA, - SENSE_DEVICE_UPDATE, - SENSE_DEVICES_DATA, - SENSE_DISCOVERED_DEVICES_DATA, -) +from . import SenseConfigEntry +from .const import ATTRIBUTION, DOMAIN, MDI_ICONS, SENSE_DEVICE_UPDATE _LOGGER = logging.getLogger(__name__) async def async_setup_entry( hass: HomeAssistant, - config_entry: ConfigEntry, + config_entry: SenseConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up the Sense binary sensor.""" - data = hass.data[DOMAIN][config_entry.entry_id][SENSE_DATA] - sense_devices_data = hass.data[DOMAIN][config_entry.entry_id][SENSE_DEVICES_DATA] - sense_monitor_id = data.sense_monitor_id + sense_monitor_id = config_entry.runtime_data.data.sense_monitor_id - sense_devices = hass.data[DOMAIN][config_entry.entry_id][ - SENSE_DISCOVERED_DEVICES_DATA - ] + sense_devices = config_entry.runtime_data.discovered + device_data = config_entry.runtime_data.device_data devices = [ - SenseDevice(sense_devices_data, device, sense_monitor_id) + SenseDevice(device_data, device, sense_monitor_id) for device in sense_devices if device["tags"]["DeviceListAllowed"] == "true" ] diff --git a/homeassistant/components/sense/const.py b/homeassistant/components/sense/const.py index 3ad35ff345d..5e944c18d8d 100644 --- a/homeassistant/components/sense/const.py +++ b/homeassistant/components/sense/const.py @@ -12,11 +12,7 @@ DOMAIN = "sense" DEFAULT_TIMEOUT = 30 ACTIVE_UPDATE_RATE = 60 DEFAULT_NAME = "Sense" -SENSE_DATA = "sense_data" SENSE_DEVICE_UPDATE = "sense_devices_update" -SENSE_DEVICES_DATA = "sense_devices_data" -SENSE_DISCOVERED_DEVICES_DATA = "sense_discovered_devices" -SENSE_TRENDS_COORDINATOR = "sense_trends_coordinator" ACTIVE_NAME = "Energy" ACTIVE_TYPE = "active" diff --git a/homeassistant/components/sense/sensor.py b/homeassistant/components/sense/sensor.py index 199bae43701..129b1262fd0 100644 --- a/homeassistant/components/sense/sensor.py +++ b/homeassistant/components/sense/sensor.py @@ -5,7 +5,6 @@ from homeassistant.components.sensor import ( SensorEntity, SensorStateClass, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import ( PERCENTAGE, UnitOfElectricPotential, @@ -18,6 +17,7 @@ from homeassistant.helpers.dispatcher import async_dispatcher_connect from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.update_coordinator import CoordinatorEntity +from . import SenseConfigEntry from .const import ( ACTIVE_NAME, ACTIVE_TYPE, @@ -34,11 +34,7 @@ from .const import ( PRODUCTION_NAME, PRODUCTION_PCT_ID, PRODUCTION_PCT_NAME, - SENSE_DATA, SENSE_DEVICE_UPDATE, - SENSE_DEVICES_DATA, - SENSE_DISCOVERED_DEVICES_DATA, - SENSE_TRENDS_COORDINATOR, SOLAR_POWERED_ID, SOLAR_POWERED_NAME, TO_GRID_ID, @@ -87,26 +83,23 @@ def sense_to_mdi(sense_icon): async def async_setup_entry( hass: HomeAssistant, - config_entry: ConfigEntry, + config_entry: SenseConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up the Sense sensor.""" - base_data = hass.data[DOMAIN][config_entry.entry_id] - data = base_data[SENSE_DATA] - sense_devices_data = base_data[SENSE_DEVICES_DATA] - trends_coordinator = base_data[SENSE_TRENDS_COORDINATOR] + data = config_entry.runtime_data.data + trends_coordinator = config_entry.runtime_data.trends # Request only in case it takes longer # than 60s await trends_coordinator.async_request_refresh() sense_monitor_id = data.sense_monitor_id - sense_devices = hass.data[DOMAIN][config_entry.entry_id][ - SENSE_DISCOVERED_DEVICES_DATA - ] + sense_devices = config_entry.runtime_data.discovered + device_data = config_entry.runtime_data.device_data entities: list[SensorEntity] = [ - SenseEnergyDevice(sense_devices_data, device, sense_monitor_id) + SenseEnergyDevice(device_data, device, sense_monitor_id) for device in sense_devices if device["tags"]["DeviceListAllowed"] == "true" ]