Remove config entry from lock device when needed (#65)

Author: raman325

custom_components/lock_code_manager/__init__.py

@@ -334,6 +334,12 @@ async def async_update_listener(hass: HomeAssistant, config_entry: ConfigEntry)
             "%s (%s): Removing lock %s entities", entry_id, entry_title, lock_entity_id
         )
         async_dispatcher_send(hass, f"{DOMAIN}_{entry_id}_remove_lock", lock_entity_id)
+        lock: BaseLock = hass.data[DOMAIN][CONF_LOCKS][lock_entity_id]
+        if lock.device_entry:
+            dev_reg = dr.async_get(hass)
+            dev_reg.async_update_device(
+                lock.device_entry.id, remove_config_entry_id=entry_id
+            )
         await async_unload_lock(hass, config_entry, lock_entity_id)
 
     # Notify any existing entities that additional locks have been added then create

custom_components/lock_code_manager/entity.py

@@ -370,6 +370,7 @@ def __init__(
                 connections=lock.device_entry.connections,
                 identifiers=lock.device_entry.identifiers,
             )
+            self._attr_name = f"Code slot {slot_num}"
 
         self._attr_unique_id = (
             f"{self.base_unique_id}|{slot_num}|{self.key}|{lock.lock.entity_id}"

tests/common.py

@@ -10,6 +10,7 @@
 from homeassistant.components.lock import LockEntity
 from homeassistant.const import CONF_ENABLED, CONF_NAME, CONF_PIN
 from homeassistant.core import HomeAssistant, callback
+from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.util import slugify
 
 from custom_components.lock_code_manager.const import (
@@ -130,6 +131,10 @@ def __init__(self, name: str) -> None:
         self._attr_name = name
         self._attr_unique_id = slugify(name)
         self._attr_is_locked = False
+        self._attr_has_entity_name = False
+        self._attr_device_info = DeviceInfo(
+            identifiers={(DOMAIN, f"lock.{slugify(name)}")}, name=name
+        )
         super().__init__()
 
 

tests/test_init.py

@@ -19,7 +19,7 @@
     Platform,
 )
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import entity_registry as er
+from homeassistant.helpers import device_registry as dr, entity_registry as er
 
 from custom_components.lock_code_manager.const import (
     ATTR_PIN_SYNCED_TO_LOCKS,
@@ -49,9 +49,16 @@ async def test_entry_setup_and_unload(
     lock_code_manager_config_entry,
 ):
     """Test entry setup and unload."""
+    mock_lock_entry_id = mock_lock_config_entry.entry_id
     lcm_entry_id = lock_code_manager_config_entry.entry_id
+    dev_reg = dr.async_get(hass)
     ent_reg = er.async_get(hass)
 
+    for entity_id in (LOCK_1_ENTITY_ID, LOCK_2_ENTITY_ID):
+        device = dev_reg.async_get_device({(DOMAIN, entity_id)})
+        assert device
+        assert device.config_entries == {mock_lock_entry_id, lcm_entry_id}
+
     unique_ids = set()
     for slot in range(1, 3):
         for entity_id in (LOCK_1_ENTITY_ID, LOCK_2_ENTITY_ID):
@@ -148,6 +155,12 @@ async def test_entry_setup_and_unload(
     )
     await hass.async_block_till_done()
 
+    # Validate that the config entry is removed from the device associated with the
+    # lock that was removed from the config entry
+    device = dev_reg.async_get_device({(DOMAIN, LOCK_2_ENTITY_ID)})
+    assert device
+    assert device.config_entries == {mock_lock_entry_id}
+
     unique_ids = set()
     for slot in range(1, 3):
         unique_ids.add(f"{lcm_entry_id}|{slot}|{CONF_CODE}|{LOCK_1_ENTITY_ID}")

tests/test_sensor.py

@@ -14,9 +14,9 @@ async def test_sensor_entity(
 ):
     """Test sensor entity."""
     for code_slot, pin in ((1, "1234"), (2, "5678")):
-        state = hass.states.get(f"sensor.mock_title_code_slot_{code_slot}_code")
+        state = hass.states.get(f"sensor.test_1_code_slot_{code_slot}")
         assert state
         assert state.state == pin
-        state = hass.states.get(f"sensor.mock_title_code_slot_{code_slot}_code_2")
+        state = hass.states.get(f"sensor.test_2_code_slot_{code_slot}")
         assert state
         assert state.state == pin