Use async_get_device_and_config_entry service helper in shelly (#180162)

homeassistant/components/shelly/services.py

@@ -1,12 +1,11 @@
 """Support for services."""
 
-from typing import TYPE_CHECKING, Any, cast
+from typing import Any, cast
 
 from aioshelly.const import RPC_GENERATIONS
 from aioshelly.exceptions import DeviceConnectionError, RpcCallError
 import voluptuous as vol
 
-from homeassistant.config_entries import ConfigEntryState
 from homeassistant.const import ATTR_DEVICE_ID
 from homeassistant.core import (
     HomeAssistant,
@@ -16,7 +15,8 @@ from homeassistant.core import (
     callback,
 )
 from homeassistant.exceptions import HomeAssistantError, ServiceValidationError
-from homeassistant.helpers import config_validation as cv, device_registry as dr
+from homeassistant.helpers import config_validation as cv
+from homeassistant.helpers.service import async_get_device_and_config_entry
 from homeassistant.util.json import JsonValueType
 
 from .const import ATTR_KEY, ATTR_VALUE, CONF_SLEEP_PERIOD, DOMAIN
@@ -45,49 +45,22 @@ def async_get_config_entry_for_service_call(
     call: ServiceCall,
 ) -> ShellyConfigEntry:
     """Get the config entry related to a service call (by device ID)."""
-    device_registry = dr.async_get(call.hass)
-    device_id = call.data[ATTR_DEVICE_ID]
+    config_entry: ShellyConfigEntry
+    _, config_entry = async_get_device_and_config_entry(
+        call.hass, DOMAIN, call.data[ATTR_DEVICE_ID]
+    )
 
-    if (device_entry := device_registry.async_get(device_id)) is None:
+    if (
+        config_entry.data.get(CONF_SLEEP_PERIOD, 0) > 0
+        or get_device_entry_gen(config_entry) not in RPC_GENERATIONS
+    ):
         raise ServiceValidationError(
             translation_domain=DOMAIN,
-            translation_key="invalid_device_id",
-            translation_placeholders={"device_id": device_id},
+            translation_key="kvs_not_supported",
+            translation_placeholders={"device": config_entry.title},
         )
 
-    for entry_id in device_entry.config_entries:
-        config_entry = call.hass.config_entries.async_get_entry(entry_id)
-
-        if TYPE_CHECKING:
-            assert config_entry
-
-        if config_entry.domain != DOMAIN:
-            continue
-        if config_entry.state is not ConfigEntryState.LOADED:
-            raise ServiceValidationError(
-                translation_domain=DOMAIN,
-                translation_key="entry_not_loaded",
-                translation_placeholders={"device": config_entry.title},
-            )
-        if get_device_entry_gen(config_entry) not in RPC_GENERATIONS:
-            raise ServiceValidationError(
-                translation_domain=DOMAIN,
-                translation_key="kvs_not_supported",
-                translation_placeholders={"device": config_entry.title},
-            )
-        if config_entry.data.get(CONF_SLEEP_PERIOD, 0) > 0:
-            raise ServiceValidationError(
-                translation_domain=DOMAIN,
-                translation_key="kvs_not_supported",
-                translation_placeholders={"device": config_entry.title},
-            )
-        return config_entry
-
-    raise ServiceValidationError(
-        translation_domain=DOMAIN,
-        translation_key="config_entry_not_found",
-        translation_placeholders={"device_id": device_id},
-    )
+    return config_entry
 
 
 async def _async_execute_action(

homeassistant/components/shelly/strings.json

@@ -642,9 +642,6 @@
     "circuit_breaker_remote_disabled": {
       "message": "Circuit breaker for {entity} of {device} is tripped and cannot be turned on remotely"
     },
-    "config_entry_not_found": {
-      "message": "Config entry for device ID {device_id} not found"
-    },
     "device_communication_action_error": {
       "message": "Device communication error occurred while calling action for {entity} of {device}"
     },
@@ -657,15 +654,9 @@
     "device_not_initialized": {
       "message": "{device} not initialized"
     },
-    "entry_not_loaded": {
-      "message": "Config entry not loaded for {device}"
-    },
     "firmware_unsupported": {
       "message": "{device} is running an unsupported firmware, please update the firmware"
     },
-    "invalid_device_id": {
-      "message": "Invalid device ID specified: {device_id}"
-    },
     "invalid_trigger": {
       "message": "Invalid device automation trigger (type, subtype): {trigger}"
     },

tests/components/shelly/test_services.py

@@ -12,7 +12,7 @@ from homeassistant.components.shelly.services import (
 )
 from homeassistant.config_entries import ConfigEntryState
 from homeassistant.const import ATTR_DEVICE_ID
-from homeassistant.core import HomeAssistant
+from homeassistant.core import DOMAIN as HOMEASSISTANT_DOMAIN, HomeAssistant
 from homeassistant.exceptions import HomeAssistantError, ServiceValidationError
 from homeassistant.helpers import device_registry as dr
 
@@ -61,8 +61,8 @@ async def test_service_get_kvs_value_invalid_device(hass: HomeAssistant) -> None
             return_response=True,
         )
 
-    assert exc_info.value.translation_domain == DOMAIN
-    assert exc_info.value.translation_key == "invalid_device_id"
+    assert exc_info.value.translation_domain == HOMEASSISTANT_DOMAIN
+    assert exc_info.value.translation_key == "service_device_not_found"
     assert exc_info.value.translation_placeholders == {
         ATTR_DEVICE_ID: "invalid_device_id"
     }
@@ -149,9 +149,12 @@ async def test_config_entry_not_loaded(
             return_response=True,
         )
 
-    assert exc_info.value.translation_domain == DOMAIN
-    assert exc_info.value.translation_key == "entry_not_loaded"
-    assert exc_info.value.translation_placeholders == {"device": entry.title}
+    assert exc_info.value.translation_domain == HOMEASSISTANT_DOMAIN
+    assert exc_info.value.translation_key == "service_config_entry_not_loaded"
+    assert exc_info.value.translation_placeholders == {
+        "domain": DOMAIN,
+        "entry_title": entry.title,
+    }
 
 
 async def test_service_get_kvs_value_sleeping_device(
@@ -258,6 +261,9 @@ async def test_service_get_kvs_value_wrong_domain(
             return_response=True,
         )
 
-    assert exc_info.value.translation_domain == DOMAIN
-    assert exc_info.value.translation_key == "config_entry_not_found"
-    assert exc_info.value.translation_placeholders == {"device_id": device.id}
+    assert exc_info.value.translation_domain == HOMEASSISTANT_DOMAIN
+    assert exc_info.value.translation_key == "service_device_wrong_domain"
+    assert exc_info.value.translation_placeholders == {
+        "device_name": device.name,
+        "domain": DOMAIN,
+    }