Store runtime data inside the config entry in System Monitor (#116816)

* replace optional timeout by sane default

* SystemMonitorConfigEntry noz needed in repairs.py

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))

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(

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"

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 . import SystemMonitorConfigEntry
-from .coordinator import SystemMonitorCoordinator
 
 
 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,

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),
+            "disk_arguments": get_all_disk_mounts(hass, psutil_wrapper),
-            "network_arguments": get_all_network_interfaces(hass),
+            "network_arguments": get_all_network_interfaces(hass, psutil_wrapper),
         }
 
     cpu_temperature: float | None = None

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: