diff --git a/homeassistant/components/portainer/services.py b/homeassistant/components/portainer/services.py index 6a2530c6c9d4..556899141f96 100644 --- a/homeassistant/components/portainer/services.py +++ b/homeassistant/components/portainer/services.py @@ -51,30 +51,22 @@ SERVICE_RECREATE_CONTAINER_SCHEMA = vol.Schema( @callback -def _async_get_device(call: ServiceCall, device_id: str) -> dr.DeviceEntry: - """Get a device entry from a device ID.""" - device_reg = dr.async_get(call.hass) - if (device := device_reg.async_get(device_id, include_child_devices=False)) is None: +def _async_get_device_and_entry( + call: ServiceCall, device_id: str +) -> tuple[dr.DeviceEntry, PortainerConfigEntry]: + """Resolve and validate the device and Portainer config entry for a device ID.""" + device, config_entry = dr.async_get_device_and_config_entry_for_domain( + call.hass, device_id, domain=DOMAIN + ) + if device is None or config_entry is None: raise ServiceValidationError( translation_domain=DOMAIN, translation_key="invalid_target", ) - return device - - -@callback -def _async_get_entry_from_device( - call: ServiceCall, device: dr.DeviceEntry -) -> PortainerConfigEntry: - """Resolve and validate the Portainer config entry for a device.""" - for entry in call.hass.config_entries.async_entries(DOMAIN): - if entry.entry_id in device.config_entries: - return service.async_get_config_entry(call.hass, DOMAIN, entry.entry_id) - - raise ServiceValidationError( - translation_domain=DOMAIN, - translation_key="invalid_target", + entry: PortainerConfigEntry = service.async_get_config_entry( + call.hass, DOMAIN, config_entry.entry_id ) + return device, entry @callback @@ -122,8 +114,7 @@ def _async_get_container_and_endpoint_ids( async def prune_images(call: ServiceCall) -> None: """Prune unused images in Portainer, with more controls.""" - device = _async_get_device(call, call.data[ATTR_DEVICE_ID]) - config_entry = _async_get_entry_from_device(call, device) + device, config_entry = _async_get_device_and_entry(call, call.data[ATTR_DEVICE_ID]) coordinator = config_entry.runtime_data endpoint_id = _async_get_endpoint_id(device, config_entry) @@ -152,8 +143,9 @@ async def prune_images(call: ServiceCall) -> None: async def recreate_container(call: ServiceCall) -> None: """Recreate a container in Portainer, with more controls.""" - device = _async_get_device(call, call.data[ATTR_CONTAINER_DEVICE_ID]) - config_entry = _async_get_entry_from_device(call, device) + device, config_entry = _async_get_device_and_entry( + call, call.data[ATTR_CONTAINER_DEVICE_ID] + ) coordinator = config_entry.runtime_data endpoint_id, container_id = _async_get_container_and_endpoint_ids( device, config_entry diff --git a/tests/components/portainer/test_services.py b/tests/components/portainer/test_services.py index 5623b63fce50..f344814fbda3 100644 --- a/tests/components/portainer/test_services.py +++ b/tests/components/portainer/test_services.py @@ -20,7 +20,7 @@ from homeassistant.components.portainer.services import ( ATTR_TIMEOUT, SERVICE_PRUNE_IMAGES, SERVICE_RECREATE_CONTAINER, - _async_get_device, + _async_get_device_and_entry, ) from homeassistant.const import ATTR_DEVICE_ID from homeassistant.core import HomeAssistant @@ -297,14 +297,14 @@ async def test_service_prune_images_device_gone( mock_portainer_client: AsyncMock, mock_config_entry: MockConfigEntry, ) -> None: - """Test _async_get_device raises when the device ID no longer exists in the registry.""" + """Test resolution raises when the device ID no longer exists in the registry.""" await setup_integration(hass, mock_config_entry) mock_call = MagicMock() mock_call.hass = hass with pytest.raises(ServiceValidationError): - _async_get_device(mock_call, "nonexistent_device_id") + _async_get_device_and_entry(mock_call, "nonexistent_device_id") mock_portainer_client.images_prune.assert_not_called()