Add periodic system stats to hardware integration (#76873)

homeassistant/components/hardware/__init__.py

@@ -12,6 +12,6 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     """Set up Hardware."""
     hass.data[DOMAIN] = {}
 
-    websocket_api.async_setup(hass)
+    await websocket_api.async_setup(hass)
 
     return True

homeassistant/components/hardware/manifest.json

@@ -3,5 +3,6 @@
   "name": "Hardware",
   "config_flow": false,
   "documentation": "https://www.home-assistant.io/integrations/hardware",
-  "codeowners": ["@home-assistant/core"]
+  "codeowners": ["@home-assistant/core"],
+  "requirements": ["psutil-home-assistant==0.0.1"]
 }

homeassistant/components/hardware/websocket_api.py

@@ -2,23 +2,41 @@
 from __future__ import annotations
 
 import contextlib
-from dataclasses import asdict
+from dataclasses import asdict, dataclass
+from datetime import datetime, timedelta
 
+import psutil_home_assistant as ha_psutil
 import voluptuous as vol
 
 from homeassistant.components import websocket_api
-from homeassistant.core import HomeAssistant, callback
+from homeassistant.core import CALLBACK_TYPE, HomeAssistant, callback
 from homeassistant.exceptions import HomeAssistantError
+from homeassistant.helpers.event import async_track_time_interval
+import homeassistant.util.dt as dt_util
 
 from .const import DOMAIN
 from .hardware import async_process_hardware_platforms
 from .models import HardwareProtocol
 
 
-@callback
-def async_setup(hass: HomeAssistant) -> None:
+@dataclass
+class SystemStatus:
+    """System status."""
+
+    ha_psutil: ha_psutil
+    remove_periodic_timer: CALLBACK_TYPE | None
+    subscribers: set[tuple[websocket_api.ActiveConnection, int]]
+
+
+async def async_setup(hass: HomeAssistant) -> None:
     """Set up the hardware websocket API."""
     websocket_api.async_register_command(hass, ws_info)
+    websocket_api.async_register_command(hass, ws_subscribe_system_status)
+    hass.data[DOMAIN]["system_status"] = SystemStatus(
+        ha_psutil=await hass.async_add_executor_job(ha_psutil.PsutilWrapper),
+        remove_periodic_timer=None,
+        subscribers=set(),
+    )
 
 
 @websocket_api.websocket_command(
@@ -45,3 +63,57 @@ async def ws_info(
                 hardware_info.extend([asdict(hw) for hw in platform.async_info(hass)])
 
     connection.send_result(msg["id"], {"hardware": hardware_info})
+
+
+@websocket_api.websocket_command(
+    {
+        vol.Required("type"): "hardware/subscribe_system_status",
+    }
+)
+@websocket_api.async_response
+async def ws_subscribe_system_status(
+    hass: HomeAssistant, connection: websocket_api.ActiveConnection, msg: dict
+):
+    """Subscribe to system status updates."""
+
+    system_status: SystemStatus = hass.data[DOMAIN]["system_status"]
+
+    @callback
+    def async_update_status(now: datetime) -> None:
+        # Although cpu_percent and virtual_memory access files in the /proc vfs, those
+        # accesses do not block and we don't need to wrap the calls in an executor.
+        # https://elixir.bootlin.com/linux/v5.19.4/source/fs/proc/stat.c
+        # https://elixir.bootlin.com/linux/v5.19.4/source/fs/proc/meminfo.c#L32
+        cpu_percentage = round(
+            system_status.ha_psutil.psutil.cpu_percent(interval=None)
+        )
+        virtual_memory = system_status.ha_psutil.psutil.virtual_memory()
+        json_msg = {
+            "cpu_percent": cpu_percentage,
+            "memory_used_percent": virtual_memory.percent,
+            "memory_used_mb": round(
+                (virtual_memory.total - virtual_memory.available) / 1024**2, 1
+            ),
+            "memory_free_mb": round(virtual_memory.available / 1024**2, 1),
+            "timestamp": dt_util.utcnow().isoformat(),
+        }
+        for connection, msg_id in system_status.subscribers:
+            connection.send_message(websocket_api.event_message(msg_id, json_msg))
+
+    if not system_status.subscribers:
+        system_status.remove_periodic_timer = async_track_time_interval(
+            hass, async_update_status, timedelta(seconds=5)
+        )
+
+    system_status.subscribers.add((connection, msg["id"]))
+
+    @callback
+    def cancel_subscription() -> None:
+        system_status.subscribers.remove((connection, msg["id"]))
+        if not system_status.subscribers and system_status.remove_periodic_timer:
+            system_status.remove_periodic_timer()
+            system_status.remove_periodic_timer = None
+
+    connection.subscriptions[msg["id"]] = cancel_subscription
+
+    connection.send_message(websocket_api.result_message(msg["id"]))

requirements_all.txt

@@ -1309,6 +1309,9 @@ prometheus_client==0.7.1
 # homeassistant.components.proxmoxve
 proxmoxer==1.3.1
 
+# homeassistant.components.hardware
+psutil-home-assistant==0.0.1
+
 # homeassistant.components.systemmonitor
 psutil==5.9.1
 

requirements_test_all.txt

@@ -924,6 +924,9 @@ progettihwsw==0.1.1
 # homeassistant.components.prometheus
 prometheus_client==0.7.1
 
+# homeassistant.components.hardware
+psutil-home-assistant==0.0.1
+
 # homeassistant.components.androidtv
 pure-python-adb[async]==0.3.0.dev0
 

tests/components/hardware/test_websocket_api.py

@@ -1,7 +1,14 @@
 """Test the hardware websocket API."""
+from collections import namedtuple
+import datetime
+from unittest.mock import patch
+
+import psutil_home_assistant as ha_psutil
+
 from homeassistant.components.hardware.const import DOMAIN
 from homeassistant.core import HomeAssistant
 from homeassistant.setup import async_setup_component
+import homeassistant.util.dt as dt_util
 
 
 async def test_board_info(hass: HomeAssistant, hass_ws_client) -> None:
@@ -16,3 +23,67 @@ async def test_board_info(hass: HomeAssistant, hass_ws_client) -> None:
     assert msg["id"] == 1
     assert msg["success"]
     assert msg["result"] == {"hardware": []}
+
+
+TEST_TIME_ADVANCE_INTERVAL = datetime.timedelta(seconds=5 + 1)
+
+
+async def test_system_status_subscription(hass: HomeAssistant, hass_ws_client, freezer):
+    """Test websocket system status subscription."""
+
+    mock_psutil = None
+    orig_psutil_wrapper = ha_psutil.PsutilWrapper
+
+    def create_mock_psutil():
+        nonlocal mock_psutil
+        mock_psutil = orig_psutil_wrapper()
+        return mock_psutil
+
+    with patch(
+        "homeassistant.components.hardware.websocket_api.ha_psutil.PsutilWrapper",
+        wraps=create_mock_psutil,
+    ):
+        assert await async_setup_component(hass, DOMAIN, {})
+
+    client = await hass_ws_client(hass)
+
+    await client.send_json({"id": 1, "type": "hardware/subscribe_system_status"})
+    response = await client.receive_json()
+    assert response["success"]
+
+    VirtualMem = namedtuple("VirtualMemory", ["available", "percent", "total"])
+    vmem = VirtualMem(10 * 1024**2, 50, 30 * 1024**2)
+
+    with patch.object(
+        mock_psutil.psutil,
+        "cpu_percent",
+        return_value=123,
+    ), patch.object(
+        mock_psutil.psutil,
+        "virtual_memory",
+        return_value=vmem,
+    ):
+        freezer.tick(TEST_TIME_ADVANCE_INTERVAL)
+        await hass.async_block_till_done()
+
+    response = await client.receive_json()
+    assert response["event"] == {
+        "cpu_percent": 123,
+        "memory_free_mb": 10.0,
+        "memory_used_mb": 20.0,
+        "memory_used_percent": 50,
+        "timestamp": dt_util.utcnow().isoformat(),
+    }
+
+    # Unsubscribe
+    await client.send_json({"id": 8, "type": "unsubscribe_events", "subscription": 1})
+    response = await client.receive_json()
+    assert response["success"]
+
+    with patch.object(mock_psutil.psutil, "cpu_percent") as cpu_mock, patch.object(
+        mock_psutil.psutil, "virtual_memory"
+    ) as vmem_mock:
+        freezer.tick(TEST_TIME_ADVANCE_INTERVAL)
+        await hass.async_block_till_done()
+        cpu_mock.assert_not_called()
+        vmem_mock.assert_not_called()