Add memory/cpu percentage to Supervisor addons entities (#57468)

* Add memory/cpu percentage to Supervisor addons entities

* fixed lint

* Use a single update function. Fixed tests

* use constant

* review comments

* oops

homeassistant/components/hassio/__init__.py

@@ -1,6 +1,7 @@
 """Support for Hass.io."""
 from __future__ import annotations
 
+import asyncio
 from datetime import timedelta
 import logging
 import os
@@ -45,6 +46,7 @@ from .const import (
     ATTR_SLUG,
     ATTR_URL,
     ATTR_VERSION,
+    DATA_KEY_ADDONS,
     DOMAIN,
     SupervisorEntityModel,
 )
@@ -75,7 +77,8 @@ DATA_STORE = "hassio_store"
 DATA_INFO = "hassio_info"
 DATA_OS_INFO = "hassio_os_info"
 DATA_SUPERVISOR_INFO = "hassio_supervisor_info"
-HASSIO_UPDATE_INTERVAL = timedelta(minutes=55)
+DATA_ADDONS_STATS = "hassio_addons_stats"
+HASSIO_UPDATE_INTERVAL = timedelta(minutes=5)
 
 ADDONS_COORDINATOR = "hassio_addons_coordinator"
 
@@ -343,6 +346,16 @@ def get_supervisor_info(hass):
     return hass.data.get(DATA_SUPERVISOR_INFO)
 
 
+@callback
+@bind_hass
+def get_addons_stats(hass):
+    """Return Addons stats.
+
+    Async friendly.
+    """
+    return hass.data.get(DATA_ADDONS_STATS)
+
+
 @callback
 @bind_hass
 def get_os_info(hass):
@@ -499,25 +512,50 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:  # noqa:
             DOMAIN, service, async_service_handler, schema=settings.schema
         )
 
+    async def update_addon_stats(slug):
+        """Update single addon stats."""
+        stats = await hassio.get_addon_stats(slug)
+        return (slug, stats)
+
     async def update_info_data(now):
         """Update last available supervisor information."""
+
         try:
-            hass.data[DATA_INFO] = await hassio.get_info()
+            (
-            hass.data[DATA_HOST_INFO] = await hassio.get_host_info()
+                hass.data[DATA_INFO],
-            hass.data[DATA_STORE] = await hassio.get_store()
+                hass.data[DATA_HOST_INFO],
-            hass.data[DATA_CORE_INFO] = await hassio.get_core_info()
+                hass.data[DATA_STORE],
-            hass.data[DATA_SUPERVISOR_INFO] = await hassio.get_supervisor_info()
+                hass.data[DATA_CORE_INFO],
-            hass.data[DATA_OS_INFO] = await hassio.get_os_info()
+                hass.data[DATA_SUPERVISOR_INFO],
+                hass.data[DATA_OS_INFO],
+            ) = await asyncio.gather(
+                hassio.get_info(),
+                hassio.get_host_info(),
+                hassio.get_store(),
+                hassio.get_core_info(),
+                hassio.get_supervisor_info(),
+                hassio.get_os_info(),
+            )
+
+            addon_slugs = [
+                addon[ATTR_SLUG]
+                for addon in hass.data[DATA_SUPERVISOR_INFO].get("addons", [])
+            ]
+            stats_data = await asyncio.gather(
+                *[update_addon_stats(slug) for slug in addon_slugs]
+            )
+            hass.data[DATA_ADDONS_STATS] = dict(stats_data)
+
             if ADDONS_COORDINATOR in hass.data:
                 await hass.data[ADDONS_COORDINATOR].async_refresh()
         except HassioAPIError as err:
-            _LOGGER.warning("Can't read last version: %s", err)
+            _LOGGER.warning("Can't read Supervisor data: %s", err)
 
         hass.helpers.event.async_track_point_in_utc_time(
             update_info_data, utcnow() + HASSIO_UPDATE_INTERVAL
         )
 
-    # Fetch last version
+    # Fetch data
     await update_info_data(None)
 
     async def async_handle_core_service(call):
@@ -675,6 +713,7 @@ class HassioDataUpdateCoordinator(DataUpdateCoordinator):
         """Update data via library."""
         new_data = {}
         supervisor_info = get_supervisor_info(self.hass)
+        addons_stats = get_addons_stats(self.hass)
         store_data = get_store(self.hass)
 
         repositories = {
@@ -682,9 +721,10 @@ class HassioDataUpdateCoordinator(DataUpdateCoordinator):
             for repo in store_data.get("repositories", [])
         }
 
-        new_data["addons"] = {
+        new_data[DATA_KEY_ADDONS] = {
             addon[ATTR_SLUG]: {
                 **addon,
+                **((addons_stats or {}).get(addon[ATTR_SLUG], {})),
                 ATTR_REPOSITORY: repositories.get(
                     addon.get(ATTR_REPOSITORY), addon.get(ATTR_REPOSITORY, "")
                 ),
@@ -697,7 +737,7 @@ class HassioDataUpdateCoordinator(DataUpdateCoordinator):
         # If this is the initial refresh, register all addons and return the dict
         if not self.data:
             async_register_addons_in_dev_reg(
-                self.entry_id, self.dev_reg, new_data["addons"].values()
+                self.entry_id, self.dev_reg, new_data[DATA_KEY_ADDONS].values()
             )
             if self.is_hass_os:
                 async_register_os_in_dev_reg(
@@ -711,13 +751,15 @@ class HassioDataUpdateCoordinator(DataUpdateCoordinator):
             if self.entry_id in device.config_entries
             and device.model == SupervisorEntityModel.ADDON
         }
-        if stale_addons := supervisor_addon_devices - set(new_data["addons"]):
+        if stale_addons := supervisor_addon_devices - set(new_data[DATA_KEY_ADDONS]):
             async_remove_addons_from_dev_reg(self.dev_reg, stale_addons)
 
         # If there are new add-ons, we should reload the config entry so we can
         # create new devices and entities. We can return an empty dict because
         # coordinator will be recreated.
-        if self.data and set(new_data["addons"]) - set(self.data["addons"]):
+        if self.data and set(new_data[DATA_KEY_ADDONS]) - set(
+            self.data[DATA_KEY_ADDONS]
+        ):
             self.hass.async_create_task(
                 self.hass.config_entries.async_reload(self.entry_id)
             )

homeassistant/components/hassio/const.py

@@ -42,6 +42,8 @@ EVENT_SUPERVISOR_EVENT = "supervisor_event"
 ATTR_VERSION = "version"
 ATTR_VERSION_LATEST = "version_latest"
 ATTR_UPDATE_AVAILABLE = "update_available"
+ATTR_CPU_PERCENT = "cpu_percent"
+ATTR_MEMORY_PERCENT = "memory_percent"
 ATTR_SLUG = "slug"
 ATTR_URL = "url"
 ATTR_REPOSITORY = "repository"

homeassistant/components/hassio/handler.py

@@ -118,6 +118,14 @@ class HassIO:
         """
         return self.send_command(f"/addons/{addon}/info", method="get")
 
+    @api_data
+    def get_addon_stats(self, addon):
+        """Return stats for an Add-on.
+
+        This method returns a coroutine.
+        """
+        return self.send_command(f"/addons/{addon}/stats", method="get")
+
     @api_data
     def get_store(self):
         """Return data from the store.

homeassistant/components/hassio/sensor.py

@@ -1,16 +1,28 @@
 """Sensor platform for Hass.io addons."""
 from __future__ import annotations
 
-from homeassistant.components.sensor import SensorEntity, SensorEntityDescription
+from homeassistant.components.sensor import (
+    STATE_CLASS_MEASUREMENT,
+    SensorEntity,
+    SensorEntityDescription,
+)
 from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import PERCENTAGE
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from . import ADDONS_COORDINATOR
-from .const import ATTR_VERSION, ATTR_VERSION_LATEST, DATA_KEY_ADDONS, DATA_KEY_OS
+from .const import (
+    ATTR_CPU_PERCENT,
+    ATTR_MEMORY_PERCENT,
+    ATTR_VERSION,
+    ATTR_VERSION_LATEST,
+    DATA_KEY_ADDONS,
+    DATA_KEY_OS,
+)
 from .entity import HassioAddonEntity, HassioOSEntity
 
-ENTITY_DESCRIPTIONS = (
+COMMON_ENTITY_DESCRIPTIONS = (
     SensorEntityDescription(
         entity_registry_enabled_default=False,
         key=ATTR_VERSION,
@@ -23,6 +35,27 @@ ENTITY_DESCRIPTIONS = (
     ),
 )
 
+ADDON_ENTITY_DESCRIPTIONS = COMMON_ENTITY_DESCRIPTIONS + (
+    SensorEntityDescription(
+        entity_registry_enabled_default=False,
+        key=ATTR_CPU_PERCENT,
+        name="CPU Percent",
+        icon="mdi:cpu-64-bit",
+        native_unit_of_measurement=PERCENTAGE,
+        state_class=STATE_CLASS_MEASUREMENT,
+    ),
+    SensorEntityDescription(
+        entity_registry_enabled_default=False,
+        key=ATTR_MEMORY_PERCENT,
+        name="Memory Percent",
+        icon="mdi:memory",
+        native_unit_of_measurement=PERCENTAGE,
+        state_class=STATE_CLASS_MEASUREMENT,
+    ),
+)
+
+OS_ENTITY_DESCRIPTIONS = COMMON_ENTITY_DESCRIPTIONS
+
 
 async def async_setup_entry(
     hass: HomeAssistant,
@@ -34,8 +67,8 @@ async def async_setup_entry(
 
     entities = []
 
-    for entity_description in ENTITY_DESCRIPTIONS:
+    for addon in coordinator.data[DATA_KEY_ADDONS].values():
-        for addon in coordinator.data[DATA_KEY_ADDONS].values():
+        for entity_description in ADDON_ENTITY_DESCRIPTIONS:
             entities.append(
                 HassioAddonSensor(
                     addon=addon,
@@ -44,7 +77,8 @@ async def async_setup_entry(
                 )
             )
 
-        if coordinator.is_hass_os:
+    if coordinator.is_hass_os:
+        for entity_description in OS_ENTITY_DESCRIPTIONS:
             entities.append(
                 HassioOSSensor(
                     coordinator=coordinator,

tests/components/hassio/__init__.py

@@ -1,48 +1,2 @@
 """Tests for Hass.io component."""
-import pytest
-
 HASSIO_TOKEN = "123456"
-
-
-@pytest.fixture(autouse=True)
-def mock_all(aioclient_mock):
-    """Mock all setup requests."""
-    aioclient_mock.post("http://127.0.0.1/homeassistant/options", json={"result": "ok"})
-    aioclient_mock.get("http://127.0.0.1/supervisor/ping", json={"result": "ok"})
-    aioclient_mock.post("http://127.0.0.1/supervisor/options", json={"result": "ok"})
-    aioclient_mock.get(
-        "http://127.0.0.1/info",
-        json={
-            "result": "ok",
-            "data": {"supervisor": "222", "homeassistant": "0.110.0", "hassos": None},
-        },
-    )
-    aioclient_mock.get(
-        "http://127.0.0.1/host/info",
-        json={
-            "result": "ok",
-            "data": {
-                "result": "ok",
-                "data": {
-                    "chassis": "vm",
-                    "operating_system": "Debian GNU/Linux 10 (buster)",
-                    "kernel": "4.19.0-6-amd64",
-                },
-            },
-        },
-    )
-    aioclient_mock.get(
-        "http://127.0.0.1/core/info",
-        json={"result": "ok", "data": {"version_latest": "1.0.0"}},
-    )
-    aioclient_mock.get(
-        "http://127.0.0.1/os/info",
-        json={"result": "ok", "data": {"version_latest": "1.0.0"}},
-    )
-    aioclient_mock.get(
-        "http://127.0.0.1/supervisor/info",
-        json={"result": "ok", "data": {"version_latest": "1.0.0"}},
-    )
-    aioclient_mock.get(
-        "http://127.0.0.1/ingress/panels", json={"result": "ok", "data": {"panels": {}}}
-    )

tests/components/hassio/test_handler.py

@@ -177,6 +177,18 @@ async def test_api_addon_info(hassio_handler, aioclient_mock):
     assert aioclient_mock.call_count == 1
 
 
+async def test_api_addon_stats(hassio_handler, aioclient_mock):
+    """Test setup with API Add-on stats."""
+    aioclient_mock.get(
+        "http://127.0.0.1/addons/test/stats",
+        json={"result": "ok", "data": {"memory_percent": 0.01}},
+    )
+
+    data = await hassio_handler.get_addon_stats("test")
+    assert data["memory_percent"] == 0.01
+    assert aioclient_mock.call_count == 1
+
+
 async def test_api_discovery_message(hassio_handler, aioclient_mock):
     """Test setup with API discovery message."""
     aioclient_mock.get(

tests/components/hassio/test_init.py

@@ -90,6 +90,54 @@ def mock_all(aioclient_mock, request):
             ],
         },
     )
+    aioclient_mock.get(
+        "http://127.0.0.1/addons/test/stats",
+        json={
+            "result": "ok",
+            "data": {
+                "cpu_percent": 0.99,
+                "memory_usage": 182611968,
+                "memory_limit": 3977146368,
+                "memory_percent": 4.59,
+                "network_rx": 362570232,
+                "network_tx": 82374138,
+                "blk_read": 46010945536,
+                "blk_write": 15051526144,
+            },
+        },
+    )
+    aioclient_mock.get(
+        "http://127.0.0.1/addons/test2/stats",
+        json={
+            "result": "ok",
+            "data": {
+                "cpu_percent": 0.8,
+                "memory_usage": 51941376,
+                "memory_limit": 3977146368,
+                "memory_percent": 1.31,
+                "network_rx": 31338284,
+                "network_tx": 15692900,
+                "blk_read": 740077568,
+                "blk_write": 6004736,
+            },
+        },
+    )
+    aioclient_mock.get(
+        "http://127.0.0.1/addons/test3/stats",
+        json={
+            "result": "ok",
+            "data": {
+                "cpu_percent": 0.8,
+                "memory_usage": 51941376,
+                "memory_limit": 3977146368,
+                "memory_percent": 1.31,
+                "network_rx": 31338284,
+                "network_tx": 15692900,
+                "blk_read": 740077568,
+                "blk_write": 6004736,
+            },
+        },
+    )
     aioclient_mock.get(
         "http://127.0.0.1/ingress/panels", json={"result": "ok", "data": {"panels": {}}}
     )

tests/components/hassio/test_sensor.py

@@ -0,0 +1,172 @@
+"""The tests for the hassio sensors."""
+
+import os
+from unittest.mock import patch
+
+import pytest
+
+from homeassistant.components.hassio import DOMAIN
+from homeassistant.helpers import entity_registry
+from homeassistant.setup import async_setup_component
+
+from tests.common import MockConfigEntry
+
+MOCK_ENVIRON = {"HASSIO": "127.0.0.1", "HASSIO_TOKEN": "abcdefgh"}
+
+
+@pytest.fixture(autouse=True)
+def mock_all(aioclient_mock, request):
+    """Mock all setup requests."""
+    aioclient_mock.post("http://127.0.0.1/homeassistant/options", json={"result": "ok"})
+    aioclient_mock.get("http://127.0.0.1/supervisor/ping", json={"result": "ok"})
+    aioclient_mock.post("http://127.0.0.1/supervisor/options", json={"result": "ok"})
+    aioclient_mock.get(
+        "http://127.0.0.1/info",
+        json={
+            "result": "ok",
+            "data": {"supervisor": "222", "homeassistant": "0.110.0", "hassos": None},
+        },
+    )
+    aioclient_mock.get(
+        "http://127.0.0.1/store",
+        json={
+            "result": "ok",
+            "data": {"addons": [], "repositories": []},
+        },
+    )
+    aioclient_mock.get(
+        "http://127.0.0.1/host/info",
+        json={
+            "result": "ok",
+            "data": {
+                "result": "ok",
+                "data": {
+                    "chassis": "vm",
+                    "operating_system": "Debian GNU/Linux 10 (buster)",
+                    "kernel": "4.19.0-6-amd64",
+                },
+            },
+        },
+    )
+    aioclient_mock.get(
+        "http://127.0.0.1/core/info",
+        json={"result": "ok", "data": {"version_latest": "1.0.0", "version": "1.0.0"}},
+    )
+    aioclient_mock.get(
+        "http://127.0.0.1/os/info",
+        json={"result": "ok", "data": {"version_latest": "1.0.0", "version": "1.0.0"}},
+    )
+    aioclient_mock.get(
+        "http://127.0.0.1/supervisor/info",
+        json={
+            "result": "ok",
+            "data": {
+                "result": "ok",
+                "version_latest": "1.0.0",
+                "addons": [
+                    {
+                        "name": "test",
+                        "slug": "test",
+                        "installed": True,
+                        "update_available": False,
+                        "version": "2.0.0",
+                        "version_latest": "2.0.1",
+                        "repository": "core",
+                        "url": "https://github.com/home-assistant/addons/test",
+                    },
+                    {
+                        "name": "test2",
+                        "slug": "test2",
+                        "installed": True,
+                        "update_available": False,
+                        "version": "3.1.0",
+                        "version_latest": "3.2.0",
+                        "repository": "core",
+                        "url": "https://github.com",
+                    },
+                ],
+            },
+        },
+    )
+    aioclient_mock.get(
+        "http://127.0.0.1/addons/test/stats",
+        json={
+            "result": "ok",
+            "data": {
+                "cpu_percent": 0.99,
+                "memory_usage": 182611968,
+                "memory_limit": 3977146368,
+                "memory_percent": 4.59,
+                "network_rx": 362570232,
+                "network_tx": 82374138,
+                "blk_read": 46010945536,
+                "blk_write": 15051526144,
+            },
+        },
+    )
+    aioclient_mock.get(
+        "http://127.0.0.1/addons/test2/stats",
+        json={
+            "result": "ok",
+            "data": {
+                "cpu_percent": 0.8,
+                "memory_usage": 51941376,
+                "memory_limit": 3977146368,
+                "memory_percent": 1.31,
+                "network_rx": 31338284,
+                "network_tx": 15692900,
+                "blk_read": 740077568,
+                "blk_write": 6004736,
+            },
+        },
+    )
+    aioclient_mock.get(
+        "http://127.0.0.1/ingress/panels", json={"result": "ok", "data": {"panels": {}}}
+    )
+
+
+async def test_sensors(hass, aioclient_mock):
+    """Test hassio OS and addons sensors."""
+
+    with patch.dict(os.environ, MOCK_ENVIRON):
+        result = await async_setup_component(
+            hass,
+            "hassio",
+            {"http": {"server_port": 9999, "server_host": "127.0.0.1"}, "hassio": {}},
+        )
+        assert result
+
+    config_entry = MockConfigEntry(domain=DOMAIN, data={}, unique_id=DOMAIN)
+    config_entry.add_to_hass(hass)
+
+    await hass.async_block_till_done()
+
+    sensors = {
+        "sensor.home_assistant_operating_system_version": "1.0.0",
+        "sensor.home_assistant_operating_system_newest_version": "1.0.0",
+        "sensor.test_version": "2.0.0",
+        "sensor.test_newest_version": "2.0.1",
+        "sensor.test2_version": "3.1.0",
+        "sensor.test2_newest_version": "3.2.0",
+        "sensor.test_cpu_percent": "0.99",
+        "sensor.test2_cpu_percent": "0.8",
+        "sensor.test_memory_percent": "4.59",
+        "sensor.test2_memory_percent": "1.31",
+    }
+
+    """Check that entities are disabled by default."""
+    for sensor in sensors:
+        assert hass.states.get(sensor) is None
+
+    """Enable sensors."""
+    ent_reg = entity_registry.async_get(hass)
+    for sensor in sensors:
+        ent_reg.async_update_entity(sensor, disabled_by=None)
+    await hass.config_entries.async_reload(config_entry.entry_id)
+
+    await hass.async_block_till_done()
+
+    """Check sensor values."""
+    for sensor, value in sensors.items():
+        state = hass.states.get(sensor)
+        assert state.state == value

tests/components/hassio/test_websocket_api.py

@@ -1,4 +1,6 @@
 """Test websocket API."""
+import pytest
+
 from homeassistant.components.hassio.const import (
     ATTR_DATA,
     ATTR_ENDPOINT,
@@ -14,11 +16,53 @@ from homeassistant.core import HomeAssistant
 from homeassistant.helpers.dispatcher import async_dispatcher_send
 from homeassistant.setup import async_setup_component
 
-from . import mock_all  # noqa: F401
-
 from tests.common import async_mock_signal
 
 
+@pytest.fixture(autouse=True)
+def mock_all(aioclient_mock):
+    """Mock all setup requests."""
+    aioclient_mock.post("http://127.0.0.1/homeassistant/options", json={"result": "ok"})
+    aioclient_mock.get("http://127.0.0.1/supervisor/ping", json={"result": "ok"})
+    aioclient_mock.post("http://127.0.0.1/supervisor/options", json={"result": "ok"})
+    aioclient_mock.get(
+        "http://127.0.0.1/info",
+        json={
+            "result": "ok",
+            "data": {"supervisor": "222", "homeassistant": "0.110.0", "hassos": None},
+        },
+    )
+    aioclient_mock.get(
+        "http://127.0.0.1/host/info",
+        json={
+            "result": "ok",
+            "data": {
+                "result": "ok",
+                "data": {
+                    "chassis": "vm",
+                    "operating_system": "Debian GNU/Linux 10 (buster)",
+                    "kernel": "4.19.0-6-amd64",
+                },
+            },
+        },
+    )
+    aioclient_mock.get(
+        "http://127.0.0.1/core/info",
+        json={"result": "ok", "data": {"version_latest": "1.0.0"}},
+    )
+    aioclient_mock.get(
+        "http://127.0.0.1/os/info",
+        json={"result": "ok", "data": {"version_latest": "1.0.0"}},
+    )
+    aioclient_mock.get(
+        "http://127.0.0.1/supervisor/info",
+        json={"result": "ok", "data": {"version_latest": "1.0.0"}},
+    )
+    aioclient_mock.get(
+        "http://127.0.0.1/ingress/panels", json={"result": "ok", "data": {"panels": {}}}
+    )
+
+
 async def test_ws_subscription(hassio_env, hass: HomeAssistant, hass_ws_client):
     """Test websocket subscription."""
     assert await async_setup_component(hass, "hassio", {})