From 59349d28597c7c21bdfcdc3d22ba5a609edf4674 Mon Sep 17 00:00:00 2001 From: Michael <35783820+mib1185@users.noreply.github.com> Date: Sun, 5 May 2024 02:03:57 +0200 Subject: [PATCH] Store runtime data inside the config entry in System Monitor (#116816) * replace optional timeout by sane default * SystemMonitorConfigEntry noz needed in repairs.py --- .../components/systemmonitor/__init__.py | 23 +++++++++++++++---- .../components/systemmonitor/binary_sensor.py | 10 ++++---- .../components/systemmonitor/const.py | 1 - .../components/systemmonitor/diagnostics.py | 8 +++---- .../components/systemmonitor/sensor.py | 16 ++++++++----- .../components/systemmonitor/util.py | 14 ++++++----- 6 files changed, 45 insertions(+), 27 deletions(-) diff --git a/homeassistant/components/systemmonitor/__init__.py b/homeassistant/components/systemmonitor/__init__.py index 25c131e547c..a0053fb4953 100644 --- a/homeassistant/components/systemmonitor/__init__.py +++ b/homeassistant/components/systemmonitor/__init__.py @@ -1,5 +1,6 @@ """The System Monitor integration.""" +from dataclasses import dataclass import logging import psutil_home_assistant as ha_psutil @@ -10,7 +11,6 @@ from homeassistant.config_entries import ConfigEntry from homeassistant.const import Platform from homeassistant.core import HomeAssistant -from .const import DOMAIN, DOMAIN_COORDINATOR from .coordinator import SystemMonitorCoordinator from .util import get_all_disk_mounts @@ -18,13 +18,26 @@ _LOGGER = logging.getLogger(__name__) PLATFORMS = [Platform.BINARY_SENSOR, Platform.SENSOR] +SystemMonitorConfigEntry = ConfigEntry["SystemMonitorData"] -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: + +@dataclass +class SystemMonitorData: + """Runtime data definition.""" + + coordinator: SystemMonitorCoordinator + psutil_wrapper: ha_psutil.PsutilWrapper + + +async def async_setup_entry( + hass: HomeAssistant, entry: SystemMonitorConfigEntry +) -> bool: """Set up System Monitor from a config entry.""" psutil_wrapper = await hass.async_add_executor_job(ha_psutil.PsutilWrapper) - hass.data[DOMAIN] = psutil_wrapper - disk_arguments = list(await hass.async_add_executor_job(get_all_disk_mounts, hass)) + disk_arguments = list( + await hass.async_add_executor_job(get_all_disk_mounts, hass, psutil_wrapper) + ) legacy_resources: set[str] = set(entry.options.get("resources", [])) for resource in legacy_resources: if resource.startswith("disk_"): @@ -40,7 +53,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: hass, psutil_wrapper, disk_arguments ) await coordinator.async_config_entry_first_refresh() - hass.data[DOMAIN_COORDINATOR] = coordinator + entry.runtime_data = SystemMonitorData(coordinator, psutil_wrapper) await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) entry.async_on_unload(entry.add_update_listener(update_listener)) diff --git a/homeassistant/components/systemmonitor/binary_sensor.py b/homeassistant/components/systemmonitor/binary_sensor.py index 9efd6f3b4e0..157ec54920b 100644 --- a/homeassistant/components/systemmonitor/binary_sensor.py +++ b/homeassistant/components/systemmonitor/binary_sensor.py @@ -17,7 +17,6 @@ from homeassistant.components.binary_sensor import ( BinarySensorEntity, BinarySensorEntityDescription, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import EntityCategory from homeassistant.core import HomeAssistant from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo @@ -25,7 +24,8 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.update_coordinator import CoordinatorEntity from homeassistant.util import slugify -from .const import CONF_PROCESS, DOMAIN, DOMAIN_COORDINATOR +from . import SystemMonitorConfigEntry +from .const import CONF_PROCESS, DOMAIN from .coordinator import SystemMonitorCoordinator _LOGGER = logging.getLogger(__name__) @@ -89,10 +89,12 @@ SENSOR_TYPES: tuple[SysMonitorBinarySensorEntityDescription, ...] = ( async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, + entry: SystemMonitorConfigEntry, + async_add_entities: AddEntitiesCallback, ) -> None: """Set up System Montor binary sensors based on a config entry.""" - coordinator: SystemMonitorCoordinator = hass.data[DOMAIN_COORDINATOR] + coordinator = entry.runtime_data.coordinator async_add_entities( SystemMonitorSensor( diff --git a/homeassistant/components/systemmonitor/const.py b/homeassistant/components/systemmonitor/const.py index 4a6000323d5..798cb82f8ef 100644 --- a/homeassistant/components/systemmonitor/const.py +++ b/homeassistant/components/systemmonitor/const.py @@ -1,7 +1,6 @@ """Constants for System Monitor.""" DOMAIN = "systemmonitor" -DOMAIN_COORDINATOR = "systemmonitor_coordinator" CONF_INDEX = "index" CONF_PROCESS = "process" diff --git a/homeassistant/components/systemmonitor/diagnostics.py b/homeassistant/components/systemmonitor/diagnostics.py index 317758651d7..7a81f1598ea 100644 --- a/homeassistant/components/systemmonitor/diagnostics.py +++ b/homeassistant/components/systemmonitor/diagnostics.py @@ -4,18 +4,16 @@ from __future__ import annotations from typing import Any -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant -from .const import DOMAIN_COORDINATOR -from .coordinator import SystemMonitorCoordinator +from . import SystemMonitorConfigEntry async def async_get_config_entry_diagnostics( - hass: HomeAssistant, entry: ConfigEntry + hass: HomeAssistant, entry: SystemMonitorConfigEntry ) -> dict[str, Any]: """Return diagnostics for Sensibo config entry.""" - coordinator: SystemMonitorCoordinator = hass.data[DOMAIN_COORDINATOR] + coordinator = entry.runtime_data.coordinator diag_data = { "last_update_success": coordinator.last_update_success, diff --git a/homeassistant/components/systemmonitor/sensor.py b/homeassistant/components/systemmonitor/sensor.py index e20f4703ab8..947f637c572 100644 --- a/homeassistant/components/systemmonitor/sensor.py +++ b/homeassistant/components/systemmonitor/sensor.py @@ -25,7 +25,7 @@ from homeassistant.components.sensor import ( SensorEntityDescription, SensorStateClass, ) -from homeassistant.config_entries import SOURCE_IMPORT, ConfigEntry +from homeassistant.config_entries import SOURCE_IMPORT from homeassistant.const import ( CONF_RESOURCES, CONF_TYPE, @@ -47,7 +47,8 @@ from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType, StateTyp from homeassistant.helpers.update_coordinator import CoordinatorEntity from homeassistant.util import slugify -from .const import CONF_PROCESS, DOMAIN, DOMAIN_COORDINATOR, NET_IO_TYPES +from . import SystemMonitorConfigEntry +from .const import CONF_PROCESS, DOMAIN, NET_IO_TYPES from .coordinator import SystemMonitorCoordinator from .util import get_all_disk_mounts, get_all_network_interfaces, read_cpu_temperature @@ -501,20 +502,23 @@ async def async_setup_platform( async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, + entry: SystemMonitorConfigEntry, + async_add_entities: AddEntitiesCallback, ) -> None: """Set up System Montor sensors based on a config entry.""" entities: list[SystemMonitorSensor] = [] legacy_resources: set[str] = set(entry.options.get("resources", [])) loaded_resources: set[str] = set() - coordinator: SystemMonitorCoordinator = hass.data[DOMAIN_COORDINATOR] + coordinator = entry.runtime_data.coordinator + psutil_wrapper = entry.runtime_data.psutil_wrapper sensor_data = coordinator.data def get_arguments() -> dict[str, Any]: """Return startup information.""" return { - "disk_arguments": get_all_disk_mounts(hass), - "network_arguments": get_all_network_interfaces(hass), + "disk_arguments": get_all_disk_mounts(hass, psutil_wrapper), + "network_arguments": get_all_network_interfaces(hass, psutil_wrapper), } cpu_temperature: float | None = None diff --git a/homeassistant/components/systemmonitor/util.py b/homeassistant/components/systemmonitor/util.py index 1889e443b2d..2a4b889bdde 100644 --- a/homeassistant/components/systemmonitor/util.py +++ b/homeassistant/components/systemmonitor/util.py @@ -8,16 +8,17 @@ import psutil_home_assistant as ha_psutil from homeassistant.core import HomeAssistant -from .const import CPU_SENSOR_PREFIXES, DOMAIN +from .const import CPU_SENSOR_PREFIXES _LOGGER = logging.getLogger(__name__) SKIP_DISK_TYPES = {"proc", "tmpfs", "devtmpfs"} -def get_all_disk_mounts(hass: HomeAssistant) -> set[str]: +def get_all_disk_mounts( + hass: HomeAssistant, psutil_wrapper: ha_psutil.PsutilWrapper +) -> set[str]: """Return all disk mount points on system.""" - psutil_wrapper: ha_psutil = hass.data[DOMAIN] disks: set[str] = set() for part in psutil_wrapper.psutil.disk_partitions(all=True): if os.name == "nt": @@ -53,9 +54,10 @@ def get_all_disk_mounts(hass: HomeAssistant) -> set[str]: return disks -def get_all_network_interfaces(hass: HomeAssistant) -> set[str]: +def get_all_network_interfaces( + hass: HomeAssistant, psutil_wrapper: ha_psutil.PsutilWrapper +) -> set[str]: """Return all network interfaces on system.""" - psutil_wrapper: ha_psutil = hass.data[DOMAIN] interfaces: set[str] = set() for interface in psutil_wrapper.psutil.net_if_addrs(): if interface.startswith("veth"): @@ -68,7 +70,7 @@ def get_all_network_interfaces(hass: HomeAssistant) -> set[str]: def get_all_running_processes(hass: HomeAssistant) -> set[str]: """Return all running processes on system.""" - psutil_wrapper: ha_psutil = hass.data.get(DOMAIN, ha_psutil.PsutilWrapper()) + psutil_wrapper = ha_psutil.PsutilWrapper() processes: set[str] = set() for proc in psutil_wrapper.psutil.process_iter(["name"]): if proc.name() not in processes: