Add prune images service to Portainer (#161009)

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Joostlek <joostlek@outlook.com>

homeassistant/components/portainer/__init__.py

@@ -15,8 +15,12 @@ from homeassistant.const import (
 )
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.aiohttp_client import async_create_clientsession
+import homeassistant.helpers.config_validation as cv
+from homeassistant.helpers.typing import ConfigType
 
+from .const import DOMAIN
 from .coordinator import PortainerCoordinator
+from .services import async_setup_services
 
 _PLATFORMS: list[Platform] = [
     Platform.BINARY_SENSOR,
@@ -25,6 +29,7 @@ _PLATFORMS: list[Platform] = [
     Platform.BUTTON,
 ]
 
+CONFIG_SCHEMA = cv.config_entry_only_config_schema(DOMAIN)
 
 type PortainerConfigEntry = ConfigEntry[PortainerCoordinator]
 
@@ -49,6 +54,12 @@ async def async_setup_entry(hass: HomeAssistant, entry: PortainerConfigEntry) ->
     return True
 
 
+async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
+    """Set up the Portainer integration."""
+    await async_setup_services(hass)
+    return True
+
+
 async def async_unload_entry(hass: HomeAssistant, entry: PortainerConfigEntry) -> bool:
     """Unload a config entry."""
     return await hass.config_entries.async_unload_platforms(entry, _PLATFORMS)

homeassistant/components/portainer/const.py

@@ -3,6 +3,7 @@
 DOMAIN = "portainer"
 DEFAULT_NAME = "Portainer"
 
+
 ENDPOINT_STATUS_DOWN = 2
 
 CONTAINER_STATE_RUNNING = "running"

homeassistant/components/portainer/icons.json

@@ -67,5 +67,10 @@
         }
       }
     }
+  },
+  "services": {
+    "prune_images": {
+      "service": "mdi:delete-sweep"
+    }
   }
 }

homeassistant/components/portainer/quality_scale.yaml

@@ -7,10 +7,7 @@ rules:
   config-flow-test-coverage: done
   config-flow: done
   dependency-transparency: done
-  docs-actions:
-    status: exempt
-    comment: |
-      No custom actions are defined.
+  docs-actions: done
   docs-high-level-description: done
   docs-installation-instructions: done
   docs-removal-instructions: done
@@ -33,10 +30,7 @@ rules:
   entity-unavailable: done
   integration-owner: done
   log-when-unavailable: done
-  parallel-updates:
-    status: exempt
-    comment: |
-      No explicit parallel updates are defined.
+  parallel-updates: todo
   reauthentication-flow:
     status: todo
     comment: |

homeassistant/components/portainer/services.py

@@ -0,0 +1,115 @@
+"""Services for the Portainer integration."""
+
+from datetime import timedelta
+
+from pyportainer import (
+    PortainerAuthenticationError,
+    PortainerConnectionError,
+    PortainerTimeoutError,
+)
+import voluptuous as vol
+
+from homeassistant.const import ATTR_DEVICE_ID
+from homeassistant.core import HomeAssistant, ServiceCall
+from homeassistant.exceptions import HomeAssistantError, ServiceValidationError
+from homeassistant.helpers import config_validation as cv, device_registry as dr
+from homeassistant.helpers.service import async_extract_config_entry_ids
+
+from .const import DOMAIN
+from .coordinator import PortainerConfigEntry
+
+ATTR_DATE_UNTIL = "until"
+ATTR_DANGLING = "dangling"
+
+SERVICE_PRUNE_IMAGES = "prune_images"
+SERVICE_PRUNE_IMAGES_SCHEMA = vol.Schema(
+    {
+        vol.Required(ATTR_DEVICE_ID): cv.string,
+        vol.Optional(ATTR_DATE_UNTIL): vol.All(
+            cv.time_period, vol.Range(min=timedelta(minutes=1))
+        ),
+        vol.Optional(ATTR_DANGLING): cv.boolean,
+    },
+)
+
+
+async def _extract_config_entry(service_call: ServiceCall) -> PortainerConfigEntry:
+    """Extract config entry from the service call."""
+    target_entry_ids = await async_extract_config_entry_ids(service_call)
+    target_entries: list[PortainerConfigEntry] = [
+        loaded_entry
+        for loaded_entry in service_call.hass.config_entries.async_loaded_entries(
+            DOMAIN
+        )
+        if loaded_entry.entry_id in target_entry_ids
+    ]
+    if not target_entries:
+        raise ServiceValidationError(
+            translation_domain=DOMAIN,
+            translation_key="invalid_target",
+        )
+    return target_entries[0]
+
+
+async def _get_endpoint_id(
+    call: ServiceCall,
+    config_entry: PortainerConfigEntry,
+) -> int:
+    """Get endpoint data from device ID."""
+    device_reg = dr.async_get(call.hass)
+    device_id = call.data[ATTR_DEVICE_ID]
+    device = device_reg.async_get(device_id)
+    assert device
+    coordinator = config_entry.runtime_data
+
+    endpoint_data = None
+    for data in coordinator.data.values():
+        if (
+            DOMAIN,
+            f"{config_entry.entry_id}_{data.endpoint.id}",
+        ) in device.identifiers:
+            endpoint_data = data
+            break
+
+    assert endpoint_data
+    return endpoint_data.endpoint.id
+
+
+async def prune_images(call: ServiceCall) -> None:
+    """Prune unused images in Portainer, with more controls."""
+    config_entry = await _extract_config_entry(call)
+    coordinator = config_entry.runtime_data
+    endpoint_id = await _get_endpoint_id(call, config_entry)
+
+    try:
+        await coordinator.portainer.images_prune(
+            endpoint_id=endpoint_id,
+            until=call.data.get(ATTR_DATE_UNTIL),
+            dangling=call.data.get(ATTR_DANGLING, False),
+        )
+    except PortainerAuthenticationError as err:
+        raise HomeAssistantError(
+            translation_domain=DOMAIN,
+            translation_key="invalid_auth_no_details",
+        ) from err
+    except PortainerConnectionError as err:
+        raise HomeAssistantError(
+            translation_domain=DOMAIN,
+            translation_key="cannot_connect_no_details",
+        ) from err
+    except PortainerTimeoutError as err:
+        raise HomeAssistantError(
+            translation_domain=DOMAIN,
+            translation_key="timeout_connect_no_details",
+        ) from err
+
+
+async def async_setup_services(hass: HomeAssistant) -> None:
+    """Set up services."""
+
+    hass.services.async_register(
+        DOMAIN,
+        SERVICE_PRUNE_IMAGES,
+        prune_images,
+        SERVICE_PRUNE_IMAGES_SCHEMA,
+    )

homeassistant/components/portainer/services.yaml

@@ -0,0 +1,18 @@
+# Services for Portainer
+
+prune_images:
+  fields:
+    device_id:
+      required: true
+      selector:
+        device:
+          integration: portainer
+          model: Endpoint
+    until:
+      required: false
+      selector:
+        duration:
+    dangling:
+      required: false
+      selector:
+        boolean: {}

homeassistant/components/portainer/strings.json

@@ -155,11 +155,34 @@
     "invalid_auth_no_details": {
       "message": "An error occurred while trying to authenticate."
     },
+    "invalid_target": {
+      "message": "Invalid device targeted."
+    },
     "timeout_connect": {
       "message": "A timeout occurred while trying to connect to the Portainer instance: {error}"
     },
     "timeout_connect_no_details": {
       "message": "A timeout occurred while trying to connect to the Portainer instance."
     }
+  },
+  "services": {
+    "prune_images": {
+      "description": "Prune unused images on a Portainer endpoint.",
+      "fields": {
+        "dangling": {
+          "description": "If true, only prune dangling images.",
+          "name": "Dangling"
+        },
+        "device_id": {
+          "description": "The endpoint to prune images on.",
+          "name": "Endpoint"
+        },
+        "until": {
+          "description": "Prune images unused for at least this time duration in the past. If not provided, all unused images will be pruned.",
+          "name": "Until"
+        }
+      },
+      "name": "Prune unused images"
+    }
   }
 }

tests/components/portainer/conftest.py

@@ -23,6 +23,8 @@ MOCK_TEST_CONFIG = {
     CONF_VERIFY_SSL: True,
 }
 
+TEST_ENTRY = "portainer_test_entry_123"
+
 
 @pytest.fixture
 def mock_setup_entry() -> Generator[AsyncMock]:
@@ -78,6 +80,6 @@ def mock_config_entry() -> MockConfigEntry:
         title="Portainer test",
         data=MOCK_TEST_CONFIG,
         unique_id=MOCK_TEST_CONFIG[CONF_API_TOKEN],
-        entry_id="portainer_test_entry_123",
+        entry_id=TEST_ENTRY,
         version=2,
     )

tests/components/portainer/test_services.py

@@ -0,0 +1,189 @@
+"""Test for Portainer services."""
+
+from datetime import timedelta
+from unittest.mock import AsyncMock
+
+from pyportainer import (
+    PortainerAuthenticationError,
+    PortainerConnectionError,
+    PortainerTimeoutError,
+)
+import pytest
+from voluptuous import MultipleInvalid
+
+from homeassistant.components.portainer.const import DOMAIN
+from homeassistant.components.portainer.services import (
+    ATTR_DANGLING,
+    ATTR_DATE_UNTIL,
+    SERVICE_PRUNE_IMAGES,
+)
+from homeassistant.const import ATTR_DEVICE_ID
+from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import HomeAssistantError, ServiceValidationError
+from homeassistant.helpers.device_registry import DeviceRegistry
+
+from . import setup_integration
+from .conftest import TEST_ENTRY
+
+from tests.common import MockConfigEntry
+
+TEST_ENDPOINT_ID = 1
+TEST_DEVICE_IDENTIFIER = f"{TEST_ENTRY}_{TEST_ENDPOINT_ID}"
+
+
+async def test_services(
+    hass: HomeAssistant,
+    device_registry: DeviceRegistry,
+    mock_portainer_client: AsyncMock,
+    mock_config_entry: MockConfigEntry,
+) -> None:
+    """Tests that the services are correct."""
+
+    await setup_integration(hass, mock_config_entry)
+    device = device_registry.async_get_device(
+        identifiers={(DOMAIN, TEST_DEVICE_IDENTIFIER)}
+    )
+    assert device is not None
+    await hass.services.async_call(
+        DOMAIN,
+        SERVICE_PRUNE_IMAGES,
+        {
+            ATTR_DEVICE_ID: device.id,
+            ATTR_DATE_UNTIL: timedelta(hours=24),
+        },
+        blocking=True,
+    )
+    mock_portainer_client.images_prune.assert_called_once_with(
+        endpoint_id=TEST_ENDPOINT_ID,
+        until=timedelta(hours=24),
+        dangling=False,
+    )
+
+
+@pytest.mark.parametrize(
+    ("call_arguments", "expected_until", "expected_dangling"),
+    [
+        ({}, None, False),
+        ({ATTR_DATE_UNTIL: timedelta(hours=12)}, timedelta(hours=12), False),
+        (
+            {ATTR_DATE_UNTIL: timedelta(hours=12), ATTR_DANGLING: True},
+            timedelta(hours=12),
+            True,
+        ),
+    ],
+    ids=["no optional", "with duration", "with duration and dangling"],
+)
+async def test_service_prune_images(
+    hass: HomeAssistant,
+    device_registry: DeviceRegistry,
+    mock_portainer_client: AsyncMock,
+    mock_config_entry: MockConfigEntry,
+    call_arguments: dict,
+    expected_until: timedelta | None,
+    expected_dangling: bool,
+) -> None:
+    """Test prune images service with the variants."""
+
+    await setup_integration(hass, mock_config_entry)
+    device = device_registry.async_get_device(
+        identifiers={(DOMAIN, TEST_DEVICE_IDENTIFIER)}
+    )
+    assert device is not None
+    await hass.services.async_call(
+        DOMAIN,
+        SERVICE_PRUNE_IMAGES,
+        {ATTR_DEVICE_ID: device.id, **call_arguments},
+        blocking=True,
+    )
+    mock_portainer_client.images_prune.assert_called_once_with(
+        endpoint_id=TEST_ENDPOINT_ID,
+        until=expected_until,
+        dangling=expected_dangling,
+    )
+
+
+async def test_service_validation_errors(
+    hass: HomeAssistant,
+    device_registry: DeviceRegistry,
+    mock_portainer_client: AsyncMock,
+    mock_config_entry: MockConfigEntry,
+) -> None:
+    """Tests that the Portainer services handle bad data."""
+
+    await setup_integration(hass, mock_config_entry)
+    device = device_registry.async_get_device(
+        identifiers={(DOMAIN, TEST_DEVICE_IDENTIFIER)}
+    )
+    assert device is not None
+
+    # Test missing device_id
+    with pytest.raises(MultipleInvalid, match="required key not provided"):
+        await hass.services.async_call(
+            DOMAIN,
+            SERVICE_PRUNE_IMAGES,
+            {},
+            blocking=True,
+        )
+    mock_portainer_client.images_prune.assert_not_called()
+
+    # Test invalid until (too short, needs to be at least 1 minute)
+    with pytest.raises(MultipleInvalid, match="value must be at least"):
+        await hass.services.async_call(
+            DOMAIN,
+            SERVICE_PRUNE_IMAGES,
+            {ATTR_DEVICE_ID: device.id, ATTR_DATE_UNTIL: timedelta(seconds=30)},
+            blocking=True,
+        )
+    mock_portainer_client.images_prune.assert_not_called()
+
+    # Test invalid device
+    with pytest.raises(ServiceValidationError, match="Invalid device targeted"):
+        await hass.services.async_call(
+            DOMAIN,
+            SERVICE_PRUNE_IMAGES,
+            {ATTR_DEVICE_ID: "invalid_device_id"},
+            blocking=True,
+        )
+    mock_portainer_client.images_prune.assert_not_called()
+
+
+@pytest.mark.parametrize(
+    ("exception", "message"),
+    [
+        (
+            PortainerAuthenticationError("auth"),
+            "An error occurred while trying to authenticate",
+        ),
+        (
+            PortainerConnectionError("conn"),
+            "An error occurred while trying to connect to the Portainer instance",
+        ),
+        (
+            PortainerTimeoutError("timeout"),
+            "A timeout occurred while trying to connect to the Portainer instance",
+        ),
+    ],
+)
+async def test_service_portainer_exceptions(
+    hass: HomeAssistant,
+    device_registry: DeviceRegistry,
+    mock_portainer_client: AsyncMock,
+    mock_config_entry: MockConfigEntry,
+    exception: HomeAssistantError,
+    message: str,
+) -> None:
+    """Test service handles Portainer exceptions."""
+    await setup_integration(hass, mock_config_entry)
+    device = device_registry.async_get_device(
+        identifiers={(DOMAIN, TEST_DEVICE_IDENTIFIER)}
+    )
+
+    mock_portainer_client.images_prune.side_effect = exception
+    with pytest.raises(HomeAssistantError, match=message):
+        await hass.services.async_call(
+            DOMAIN,
+            SERVICE_PRUNE_IMAGES,
+            {ATTR_DEVICE_ID: device.id},
+            blocking=True,
+        )
+    mock_portainer_client.images_prune.assert_called_once()