diff --git a/homeassistant/components/hardware/__init__.py b/homeassistant/components/hardware/__init__.py index b3f342d4e32..a1198534213 100644 --- a/homeassistant/components/hardware/__init__.py +++ b/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 diff --git a/homeassistant/components/hardware/manifest.json b/homeassistant/components/hardware/manifest.json index e7e156b6065..94571ce4528 100644 --- a/homeassistant/components/hardware/manifest.json +++ b/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"] } diff --git a/homeassistant/components/hardware/websocket_api.py b/homeassistant/components/hardware/websocket_api.py index df3b8868053..5c4b14570a9 100644 --- a/homeassistant/components/hardware/websocket_api.py +++ b/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"])) diff --git a/requirements_all.txt b/requirements_all.txt index 617f9615921..2f4680a361d 100644 --- a/requirements_all.txt +++ b/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 diff --git a/requirements_test_all.txt b/requirements_test_all.txt index b909571ac52..93ef4590930 100644 --- a/requirements_test_all.txt +++ b/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 diff --git a/tests/components/hardware/test_websocket_api.py b/tests/components/hardware/test_websocket_api.py index 116879aa628..bc6fb5f11dd 100644 --- a/tests/components/hardware/test_websocket_api.py +++ b/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()