Merge pull request #7 from astrayel/claude/fix-unraid-integration-bug-wIuR7

Add Docker aggregate sensors

Author: astrayel

custom_components/unraid_api/sensor.py

@@ -25,9 +25,18 @@
 from homeassistant.helpers.device_registry import DeviceInfo
 
 from . import _LOGGER
-from .const import CONF_DRIVES, CONF_SHARES, CONF_VMS, DOMAIN
+from .const import CONF_DOCKER_MODE, CONF_DRIVES, CONF_SHARES, CONF_VMS, DOCKER_MODE_OFF, DOMAIN
 from .entity import UnraidBaseEntity, UnraidEntityDescription
-from .models import Disk, DiskType, DockerContainer, Share, UpsDevice, VirtualMachine, VmState
+from .models import (
+    ContainerState,
+    Disk,
+    DiskType,
+    DockerContainer,
+    Share,
+    UpsDevice,
+    VirtualMachine,
+    VmState,
+)
 
 if TYPE_CHECKING:
     from collections.abc import Callable
@@ -406,6 +415,35 @@ def calc_disk_usage_percentage(disk: Disk) -> StateType:
     ),
 )
 
+DOCKER_AGGREGATE_SENSOR_DESCRIPTIONS: tuple[UnraidSensorEntityDescription, ...] = (
+    UnraidSensorEntityDescription(
+        key="docker_running",
+        state_class=SensorStateClass.MEASUREMENT,
+        value_fn=lambda coordinator: sum(
+            1
+            for c in coordinator.data["docker_containers"].values()
+            if c.state == ContainerState.RUNNING
+        ),
+        entity_category=EntityCategory.DIAGNOSTIC,
+    ),
+    UnraidSensorEntityDescription(
+        key="docker_exited",
+        state_class=SensorStateClass.MEASUREMENT,
+        value_fn=lambda coordinator: sum(
+            1
+            for c in coordinator.data["docker_containers"].values()
+            if c.state == ContainerState.EXITED
+        ),
+        entity_category=EntityCategory.DIAGNOSTIC,
+    ),
+    UnraidSensorEntityDescription(
+        key="docker_total",
+        state_class=SensorStateClass.MEASUREMENT,
+        value_fn=lambda coordinator: len(coordinator.data["docker_containers"]),
+        entity_category=EntityCategory.DIAGNOSTIC,
+    ),
+)
+
 VM_AGGREGATE_SENSOR_DESCRIPTIONS: tuple[UnraidSensorEntityDescription, ...] = (
     UnraidSensorEntityDescription(
         key="vms_running",
@@ -435,6 +473,12 @@ async def async_setup_entry(
         for description in SENSOR_DESCRIPTIONS
         if description.min_version <= config_entry.runtime_data.coordinator.api_client.version
     ]
+    if config_entry.options.get(CONF_DOCKER_MODE, DOCKER_MODE_OFF) != DOCKER_MODE_OFF:
+        entities.extend(
+            UnraidSensor(description, config_entry)
+            for description in DOCKER_AGGREGATE_SENSOR_DESCRIPTIONS
+            if description.min_version <= config_entry.runtime_data.coordinator.api_client.version
+        )
     if config_entry.options.get(CONF_VMS):
         entities.extend(
             UnraidSensor(description, config_entry)

custom_components/unraid_api/translations/en.json

@@ -166,6 +166,15 @@
                     "exited": "Exited"
                 }
             },
+            "docker_running": {
+                "name": "Docker containers running"
+            },
+            "docker_exited": {
+                "name": "Docker containers exited"
+            },
+            "docker_total": {
+                "name": "Docker containers total"
+            },
             "vm_state": {
                 "name": "State",
                 "state": {

tests/test_sensor.py

@@ -538,6 +538,57 @@ async def test_docker_sensors_removed(
     assert hass.states.get("sensor.test_server_grafana_public_state")
 
 
+@pytest.mark.usefixtures("entity_registry_enabled_by_default")
+async def test_docker_aggregate_sensors(hass: HomeAssistant, mock_api_client: MagicMock) -> None:
+    """Test Docker aggregate sensors on the root device."""
+    assert mock_api_client
+    assert await setup_config_entry(hass)
+
+    state = hass.states.get("sensor.test_server_docker_containers_running")
+    assert state.state == "2"
+
+    state = hass.states.get("sensor.test_server_docker_containers_exited")
+    assert state.state == "1"
+
+    state = hass.states.get("sensor.test_server_docker_containers_total")
+    assert state.state == "3"
+
+
+@pytest.mark.usefixtures("entity_registry_enabled_by_default")
+async def test_docker_aggregate_sensors_disabled(
+    hass: HomeAssistant,
+    mock_api_client: MagicMock,  # noqa: ARG001
+) -> None:
+    """Test Docker aggregate sensors are absent when Docker monitoring is off."""
+    assert await setup_config_entry(hass, options=MOCK_OPTION_DATA_DISABLED)
+
+    assert hass.states.get("sensor.test_server_docker_containers_running") is None
+    assert hass.states.get("sensor.test_server_docker_containers_exited") is None
+    assert hass.states.get("sensor.test_server_docker_containers_total") is None
+
+
+@pytest.mark.usefixtures("entity_registry_enabled_by_default")
+async def test_docker_aggregate_sensors_update(
+    hass: HomeAssistant, mock_api_client: MagicMock
+) -> None:
+    """Test Docker aggregate sensors update when containers change."""
+    api_client: MockApiClient = mock_api_client.return_value
+
+    config_entry = await setup_config_entry(hass)
+    assert config_entry
+
+    assert hass.states.get("sensor.test_server_docker_containers_running").state == "2"
+    assert hass.states.get("sensor.test_server_docker_containers_total").state == "3"
+
+    api_client.state.docker.pop(0)
+    await config_entry.runtime_data.coordinator.async_refresh()
+    await hass.async_block_till_done()
+
+    assert hass.states.get("sensor.test_server_docker_containers_running").state == "1"
+    assert hass.states.get("sensor.test_server_docker_containers_exited").state == "1"
+    assert hass.states.get("sensor.test_server_docker_containers_total").state == "2"
+
+
 @pytest.mark.usefixtures("entity_registry_enabled_by_default")
 @pytest.mark.parametrize(("api_state"), API_STATES)
 async def test_vm_sensors(