Use get_device_and_config_entry_for_domain in zwave_js (#180120)

homeassistant/components/zwave_js/api.py

@@ -487,7 +487,7 @@ def async_register_api(hass: HomeAssistant) -> None:
     websocket_api.async_register_command(hass, websocket_invoke_cc_api)
     websocket_api.async_register_command(hass, websocket_backup_nvm)
     websocket_api.async_register_command(hass, websocket_restore_nvm)
-    hass.http.register_view(FirmwareUploadView(dr.async_get(hass)))
+    hass.http.register_view(FirmwareUploadView())
 
 
 @websocket_api.require_admin
@@ -2712,18 +2712,13 @@ class FirmwareUploadView(HomeAssistantView):
     url = r"/api/zwave_js/firmware/upload/{device_id}"
     name = "api:zwave_js:firmware:upload"
 
-    def __init__(self, dev_reg: dr.DeviceRegistry) -> None:
-        """Initialize view."""
-        super().__init__()
-        self._dev_reg = dev_reg
-
     @require_admin
     async def post(self, request: web.Request, device_id: str) -> web.Response:
         """Handle upload."""
         hass = request.app[KEY_HASS]
 
         try:
-            node = async_get_node_from_device_id(hass, device_id, self._dev_reg)
+            node = async_get_node_from_device_id(hass, device_id)
         except ValueError as err:
             if "not loaded" in err.args[0]:
                 raise web_exceptions.HTTPBadRequest from err
@@ -3041,7 +3036,7 @@ async def websocket_hard_reset_controller(
     @callback
     def _handle_device_added(device: dr.DeviceEntry) -> None:
         """Handle device is added."""
-        if entry.entry_id in device.config_entries:
+        if entry.entry_id == device.config_entry_id:
             connection.send_result(msg[ID], device.id)
             async_cleanup()
 

homeassistant/components/zwave_js/device_automation_helpers.py

@@ -36,21 +36,14 @@ def generate_config_parameter_subtype(config_value: ConfigurationValue) -> str:
 @callback
 def async_bypass_dynamic_config_validation(hass: HomeAssistant, device_id: str) -> bool:
     """Return whether device's config entries are not loaded."""
-    dev_reg = dr.async_get(hass)
-    if (device := dev_reg.async_get(device_id)) is None:
-        raise ValueError(f"Device {device_id} not found")
-    entry = next(
-        (
-            config_entry
-            for config_entry in hass.config_entries.async_entries(DOMAIN)
-            if config_entry.entry_id in device.config_entries
-            and config_entry.state is ConfigEntryState.LOADED
-        ),
-        None,
+    device, config_entry = dr.async_get_device_and_config_entry_for_domain(
+        hass, device_id, domain=DOMAIN
     )
-    if not entry:
+    if device is None:
+        raise ValueError(f"Device {device_id} not found")
+    if not config_entry or config_entry.state is not ConfigEntryState.LOADED:
         return True
 
     # The driver may not be ready when the config entry is loaded.
-    client = entry.runtime_data.client
+    client = config_entry.runtime_data.client
     return client.driver is None

homeassistant/components/zwave_js/device_trigger.py

@@ -258,7 +258,7 @@ async def async_get_triggers(
     }
 
     dev_reg = dr.async_get(hass)
-    node = async_get_node_from_device_id(hass, device_id, dev_reg)
+    node = async_get_node_from_device_id(hass, device_id)
 
     if node.client.driver and node.client.driver.controller.own_node == node:
         return triggers

homeassistant/components/zwave_js/helpers.py

@@ -277,38 +277,28 @@ def get_home_and_node_id_from_device_entry(
 
 
 @callback
-def async_get_node_from_device_id(
-    hass: HomeAssistant, device_id: str, dev_reg: dr.DeviceRegistry | None = None
-) -> ZwaveNode:
+def async_get_node_from_device_id(hass: HomeAssistant, device_id: str) -> ZwaveNode:
     """Get node from a device ID.
 
     Raises ValueError if device is invalid or node can't be found.
     """
-    if not dev_reg:
-        dev_reg = dr.async_get(hass)
-
-    if not (device_entry := dev_reg.async_get(device_id, include_child_devices=False)):
+    # Use the device config entry to validate that this is a valid zwave_js device
+    # and to get the client
+    device, config_entry = cast(
+        tuple[dr.DeviceEntry | None, ZwaveJSConfigEntry | None],
+        dr.async_get_device_and_config_entry_for_domain(hass, device_id, domain=DOMAIN),
+    )
+    if device is None:
         raise ValueError(f"Device ID {device_id} is not valid")
 
-    # Use device config entry ID's to validate that this is a valid zwave_js device
-    # and to get the client
-    config_entry_ids = device_entry.config_entries
-    entry: ZwaveJSConfigEntry | None = next(
-        (
-            entry
-            for entry in hass.config_entries.async_entries(DOMAIN)
-            if entry.entry_id in config_entry_ids
-        ),
-        None,
-    )
-    if entry is None:
+    if config_entry is None:
         raise ValueError(
             f"Device {device_id} is not from an existing zwave_js config entry"
         )
-    if entry.state is not ConfigEntryState.LOADED:
+    if config_entry.state is not ConfigEntryState.LOADED:
         raise ValueError(f"Device {device_id} config entry is not loaded")
 
-    client = entry.runtime_data.client
+    client = config_entry.runtime_data.client
     driver = client.driver
 
     if driver is None:
@@ -316,7 +306,7 @@ def async_get_node_from_device_id(
 
     # Get node ID from device identifier, perform some validation, and then get the
     # node
-    identifiers = get_home_and_node_id_from_device_entry(device_entry)
+    identifiers = get_home_and_node_id_from_device_entry(device)
 
     node_id = identifiers[1] if identifiers else None
 
@@ -346,30 +336,23 @@ async def async_get_provisioning_entry_from_device_id(
 
     Raises ValueError if device is invalid
     """
-    dev_reg = dr.async_get(hass)
-
-    if not (device_entry := dev_reg.async_get(device_id)):
+    # Use the device config entry to validate that this is a valid zwave_js device
+    # and to get the client
+    device, config_entry = cast(
+        tuple[dr.DeviceEntry | None, ZwaveJSConfigEntry | None],
+        dr.async_get_device_and_config_entry_for_domain(hass, device_id, domain=DOMAIN),
+    )
+    if device is None:
         raise ValueError(f"Device ID {device_id} is not valid")
 
-    # Use device config entry ID's to validate that this is a valid zwave_js device
-    # and to get the client
-    config_entry_ids = device_entry.config_entries
-    entry: ZwaveJSConfigEntry | None = next(
-        (
-            entry
-            for entry in hass.config_entries.async_entries(DOMAIN)
-            if entry.entry_id in config_entry_ids
-        ),
-        None,
-    )
-    if entry is None:
+    if config_entry is None:
         raise ValueError(
             f"Device {device_id} is not from an existing zwave_js config entry"
         )
-    if entry.state is not ConfigEntryState.LOADED:
+    if config_entry.state is not ConfigEntryState.LOADED:
         raise ValueError(f"Device {device_id} config entry is not loaded")
 
-    client = entry.runtime_data.client
+    client = config_entry.runtime_data.client
     driver = client.driver
 
     if driver is None:
@@ -391,7 +374,6 @@ def async_get_node_from_entity_id(
     hass: HomeAssistant,
     entity_id: str,
     ent_reg: er.EntityRegistry | None = None,
-    dev_reg: dr.DeviceRegistry | None = None,
 ) -> ZwaveNode:
     """Get node from an entity ID.
 
@@ -407,7 +389,7 @@ def async_get_node_from_entity_id(
     # Assert for mypy, safe because we know that zwave_js entities are always
     # tied to a device
     assert entity_entry.device_id
-    return async_get_node_from_device_id(hass, entity_entry.device_id, dev_reg)
+    return async_get_node_from_device_id(hass, entity_entry.device_id)
 
 
 @callback
@@ -426,7 +408,7 @@ def async_get_nodes_from_area_id(
     # Add devices for all entities in an area that are Z-Wave JS entities
     nodes.update(
         {
-            async_get_node_from_device_id(hass, entity.device_id, dev_reg)
+            async_get_node_from_device_id(hass, entity.device_id)
             for entity in er.async_entries_for_area(ent_reg, area_id)
             if entity.platform == DOMAIN and entity.device_id is not None
         }
@@ -434,17 +416,14 @@ def async_get_nodes_from_area_id(
     # Add devices in an area that are Z-Wave JS devices. Child devices are skipped
     # since a child device is not a Z-Wave JS node.
     nodes.update(
-        async_get_node_from_device_id(hass, device.id, dev_reg)
+        async_get_node_from_device_id(hass, device.id)
         for device in dr.async_entries_for_area(dev_reg, area_id)
         if not isinstance(device, dr.ChildDeviceEntry)
-        and any(
-            cast(
-                ZwaveJSConfigEntry,
-                hass.config_entries.async_get_entry(config_entry_id),
-            ).domain
-            == DOMAIN
-            for config_entry_id in device.config_entries
+        and (
+            config_entry := hass.config_entries.async_get_entry(device.config_entry_id)
         )
+        is not None
+        and config_entry.domain == DOMAIN
     )
 
     return nodes
@@ -466,7 +445,7 @@ def async_get_nodes_from_targets(
     # Convert all entity IDs to nodes
     for entity_id in expand_entity_ids(hass, val.get(ATTR_ENTITY_ID, [])):
         try:
-            nodes.add(async_get_node_from_entity_id(hass, entity_id, ent_reg, dev_reg))
+            nodes.add(async_get_node_from_entity_id(hass, entity_id, ent_reg))
         except ValueError as err:
             logger.warning(err.args[0])
 
@@ -477,7 +456,7 @@ def async_get_nodes_from_targets(
     # Convert all device IDs to nodes
     for device_id in val.get(ATTR_DEVICE_ID, []):
         try:
-            nodes.add(async_get_node_from_device_id(hass, device_id, dev_reg))
+            nodes.add(async_get_node_from_device_id(hass, device_id))
         except ValueError as err:
             logger.warning(err.args[0])
 
@@ -506,11 +485,10 @@ def get_zwave_value_from_config(node: ZwaveNode, config: ConfigType) -> ZwaveVal
 
 def _zwave_js_config_entry(hass: HomeAssistant, device: dr.DeviceEntry) -> str | None:
     """Find zwave_js config entry from a device."""
-    for entry_id in device.config_entries:
-        entry = hass.config_entries.async_get_entry(entry_id)
-        if entry and entry.domain == DOMAIN:
-            return entry_id
-    return None
+    _, config_entry = dr.async_get_device_and_config_entry_for_domain(
+        hass, device.id, domain=DOMAIN
+    )
+    return config_entry.entry_id if config_entry else None
 
 
 @callback
@@ -534,7 +512,7 @@ def async_get_node_status_sensor_entity_id(
     entry = hass.config_entries.async_get_entry(entry_id)
     assert entry
     client = entry.runtime_data.client
-    node = async_get_node_from_device_id(hass, device_id, dev_reg)
+    node = async_get_node_from_device_id(hass, device_id)
     return ent_reg.async_get_entity_id(
         SENSOR_DOMAIN,
         DOMAIN,

homeassistant/components/zwave_js/services.py

@@ -974,9 +974,7 @@ class ZWaveServices:
 
         for device_id in service.data.get(ATTR_DEVICE_ID, []):
             try:
-                node = async_get_node_from_device_id(
-                    self._hass, device_id, self._dev_reg
-                )
+                node = async_get_node_from_device_id(self._hass, device_id)
             except ValueError as err:
                 _LOGGER.warning(err.args[0])
                 continue
@@ -993,9 +991,7 @@ class ZWaveServices:
                     const.DOMAIN,
                 )
                 continue
-            node = async_get_node_from_entity_id(
-                self._hass, entity_id, self._ent_reg, self._dev_reg
-            )
+            node = async_get_node_from_entity_id(self._hass, entity_id, self._ent_reg)
             if (
                 value_id := get_value_id_from_unique_id(entity_entry.unique_id)
             ) is None: