Add device_registry.async_get_device_and_config_entry_for_domain (#178991)

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

homeassistant/components/blue_current/services.py

@@ -2,7 +2,7 @@
 
 import voluptuous as vol
 
-from homeassistant.config_entries import ConfigEntry, ConfigEntryState
+from homeassistant.config_entries import ConfigEntryState
 from homeassistant.const import CONF_DEVICE_ID
 from homeassistant.core import HomeAssistant, ServiceCall, callback
 from homeassistant.exceptions import ServiceValidationError
@@ -27,37 +27,27 @@ async def start_charge_session(service_call: ServiceCall) -> None:
     charging_card_id = service_call.data[CHARGING_CARD_ID]
     device_id = service_call.data[CONF_DEVICE_ID]
 
-    # Get the device based on the given device ID.
-    device = dr.async_get(service_call.hass).devices.get(device_id)
+    device, config_entry = dr.async_get_device_and_config_entry_for_domain(
+        service_call.hass, device_id, domain=DOMAIN
+    )
 
     if device is None:
         raise ServiceValidationError(
             translation_domain=DOMAIN, translation_key="invalid_device_id"
         )
 
-    blue_current_config_entry: ConfigEntry | None = None
-
-    for config_entry_id in device.config_entries:
-        config_entry = service_call.hass.config_entries.async_get_entry(config_entry_id)
-        if not config_entry or config_entry.domain != DOMAIN:
-            # Not the blue_current config entry.
-            continue
-
-        if config_entry.state is not ConfigEntryState.LOADED:
-            raise ServiceValidationError(
-                translation_domain=DOMAIN, translation_key="config_entry_not_loaded"
-            )
-
-        blue_current_config_entry = config_entry
-        break
-
-    if not blue_current_config_entry:
+    if not config_entry:
         # The device is not connected to a valid blue_current config entry.
         raise ServiceValidationError(
             translation_domain=DOMAIN, translation_key="no_config_entry"
         )
 
-    connector = blue_current_config_entry.runtime_data
+    if config_entry.state is not ConfigEntryState.LOADED:
+        raise ServiceValidationError(
+            translation_domain=DOMAIN, translation_key="config_entry_not_loaded"
+        )
+
+    connector = config_entry.runtime_data
 
     # Get the evse_id from the identifier of the device.
     evse_id = next(

homeassistant/helpers/device_registry.py

@@ -4439,6 +4439,33 @@ def async_get_device_id_by_identifier(
     return device.id
 
 
+@callback
+def async_get_device_and_config_entry_for_domain(
+    hass: HomeAssistant, device_id: str, *, domain: str
+) -> tuple[DeviceEntry | None, ConfigEntry | None]:
+    """Get the device and the config entry of the domain owning it.
+
+    Returns (None, None) for an unknown device id or if the device is a child
+    device, and (device, None) when no config entry of the domain owns the
+    device. A returned pair is consistent: for a pre-migration composite
+    device id, the device is the domain's split device, not the composite; if
+    several splits belong to config entries of the domain, which pair is
+    returned is undefined. When no split matches the domain, the restored
+    composite is returned as the device.
+    """
+    registry = async_get(hass)
+    if (device := registry.devices.get(device_id)) is not None:
+        config_entry = hass.config_entries.async_get_entry(device.config_entry_id)
+        if config_entry is not None and config_entry.domain == domain:
+            return device, config_entry
+        return device, None
+    for split in registry.async_get_devices_for_composite_device_id(device_id):
+        config_entry = hass.config_entries.async_get_entry(split.config_entry_id)
+        if config_entry is not None and config_entry.domain == domain:
+            return split, config_entry
+    return registry.async_get(device_id, include_child_devices=False), None
+
+
 def async_setup(hass: HomeAssistant) -> None:
     """Set up device registry."""
     if DATA_REGISTRY in hass.data:

tests/helpers/test_device_registry.py

@@ -9596,6 +9596,68 @@ async def test_get_composite_splits(
     assert device_registry.devices.get_composite_splits() == {}
 
 
+async def test_async_get_device_and_config_entry_for_domain(
+    hass: HomeAssistant, device_registry: dr.DeviceRegistry
+) -> None:
+    """Test getting the device and config entry of a domain owning a device."""
+    entry = MockConfigEntry(domain="domain_a")
+    entry.add_to_hass(hass)
+    device = device_registry.async_get_or_create(
+        config_entry_id=entry.entry_id, identifiers={("domain_a", "1")}
+    )
+
+    assert dr.async_get_device_and_config_entry_for_domain(
+        hass, device.id, domain="domain_a"
+    ) == (device, entry)
+    # A domain not owning the device still gets the device
+    assert dr.async_get_device_and_config_entry_for_domain(
+        hass, device.id, domain="domain_b"
+    ) == (device, None)
+    # An unknown device id
+    assert dr.async_get_device_and_config_entry_for_domain(
+        hass, "unknown_id", domain="domain_a"
+    ) == (None, None)
+
+
+@pytest.mark.parametrize("load_registries", [False])
+async def test_async_get_device_and_config_entry_for_domain_composite(
+    hass: HomeAssistant, hass_storage: dict[str, Any]
+) -> None:
+    """Test getting the device and config entry via a composite device id."""
+    entry_a = MockConfigEntry(domain="domain_a")
+    entry_a.add_to_hass(hass)
+    entry_b = MockConfigEntry(domain="domain_b")
+    entry_b.add_to_hass(hass)
+    hass_storage[dr.STORAGE_KEY] = _composite_device_storage(entry_a, entry_b)
+
+    dr.async_setup(hass)
+    await dr.async_load(hass)
+    device_registry = dr.async_get(hass)
+
+    split_a = _get_device_for_config_entry(
+        device_registry, entry_a.entry_id, identifiers={("domain_a", "1")}
+    )
+    split_b = _get_device_for_config_entry(
+        device_registry, entry_b.entry_id, identifiers={("domain_b", "1")}
+    )
+
+    # The returned pair is consistent: the domain's split device, not the composite
+    assert dr.async_get_device_and_config_entry_for_domain(
+        hass, COMPOSITE_ID, domain="domain_a"
+    ) == (split_a, entry_a)
+    assert dr.async_get_device_and_config_entry_for_domain(
+        hass, COMPOSITE_ID, domain="domain_b"
+    ) == (split_b, entry_b)
+    # A domain owning none of the splits gets the restored composite and no entry
+    device, config_entry = dr.async_get_device_and_config_entry_for_domain(
+        hass, COMPOSITE_ID, domain="domain_c"
+    )
+    assert config_entry is None
+    assert device is not None
+    assert device.id == COMPOSITE_ID
+    assert device.config_entries == {entry_a.entry_id, entry_b.entry_id}
+
+
 @pytest.mark.parametrize("load_registries", [False])
 async def test_clear_config_entry_clears_composite_primary_config_entry(
     hass: HomeAssistant, hass_storage: dict[str, Any]