rework client side entity handling

Author: dmulcahey

zha/application/platforms/__init__.py

@@ -496,21 +496,20 @@ def __init__(
         self.PLATFORM = entity_info.platform
         self._device: WebSocketClientDevice = device
         self._entity_info: BaseEntityInfoType = entity_info
-        self._attr_enabled = self._entity_info.enabled
-        self._attr_fallback_name = self._entity_info.fallback_name
-        self._attr_translation_key = self._entity_info.translation_key
-        self._attr_entity_category = self._entity_info.entity_category
-        self._attr_entity_registry_enabled_default = (
-            self._entity_info.entity_registry_enabled_default
-        )
-        self._attr_device_class = self._entity_info.device_class
-        self._attr_state_class = self._entity_info.state_class
+        self._update_attrs_from_entity_info()
 
     @property
     def info_object(self) -> BaseEntityInfoType:
         """Return a representation of the alarm control panel."""
         return self._entity_info
 
+    @info_object.setter
+    def info_object(self, entity_info: BaseEntityInfoType) -> None:
+        """Set the entity info object."""
+        self._entity_info = entity_info
+        self._update_attrs_from_entity_info()
+        self.maybe_emit_state_changed_event()
+
     @property
     def state(self) -> dict[str, Any]:
         """Return the arguments to use in the command."""
@@ -521,6 +520,7 @@ def state(self, value: dict[str, Any]) -> None:
         """Set the state of the entity."""
         self._entity_info.state = value
         self._attr_enabled = self._entity_info.enabled
+        self.maybe_emit_state_changed_event()
 
     @property
     def group_id(self) -> int | None:
@@ -562,6 +562,18 @@ def _disable(self, future: asyncio.Future) -> None:
             self._attr_enabled = False
             self.maybe_emit_state_changed_event()
 
+    def _update_attrs_from_entity_info(self) -> None:
+        """Update the entity attributes."""
+        self._attr_enabled = self._entity_info.enabled
+        self._attr_fallback_name = self._entity_info.fallback_name
+        self._attr_translation_key = self._entity_info.translation_key
+        self._attr_entity_category = self._entity_info.entity_category
+        self._attr_entity_registry_enabled_default = (
+            self._entity_info.entity_registry_enabled_default
+        )
+        self._attr_device_class = self._entity_info.device_class
+        self._attr_state_class = self._entity_info.state_class
+
     async def async_update(self) -> None:
         """Retrieve latest state."""
         self.debug("polling current state")

zha/zigbee/device.py

@@ -1146,7 +1146,7 @@ def __init__(
         self.unique_id: str = str(extended_device_info.ieee)
         self._entities: dict[tuple[Platform, str], WebSocketClientEntity] = {}
         if self._extended_device_info.entities:
-            self._build_entities()
+            self._build_or_update_entities()
 
     @property
     def quirk_id(self) -> str | None:
@@ -1172,7 +1172,7 @@ def extended_device_info(self) -> ExtendedDeviceInfo:
     def extended_device_info(self, extended_device_info: ExtendedDeviceInfo) -> None:
         """Set extended device information."""
         self._extended_device_info = extended_device_info
-        self._build_entities()
+        self._build_or_update_entities()
 
     @property
     def gateway(self) -> WebSocketClientGateway:
@@ -1300,19 +1300,18 @@ def platform_entities(self) -> dict[tuple[Platform, str], WebSocketClientEntity]
         """Return the platform entities for this device."""
         return self._entities
 
-    def _build_entities(self):
+    def _build_or_update_entities(self):
         """Build the entities for this device or rebuild them from extended device info."""
-        self._entities.update(
-            {
-                (
-                    entity_info.platform,
-                    entity_info.unique_id,
-                ): discovery.ENTITY_INFO_CLASS_TO_WEBSOCKET_CLIENT_ENTITY_CLASS[
-                    entity_info.__class__
-                ](entity_info, self)
-                for entity_info in self._extended_device_info.entities.values()
-            }
-        )
+        for entity_info in self._extended_device_info.entities.values():
+            entity_key = (entity_info.platform, entity_info.unique_id)
+            if entity_key in self._entities:
+                self._entities[entity_key].entity_info = entity_info
+            else:
+                self._entities[entity_key] = (
+                    discovery.ENTITY_INFO_CLASS_TO_WEBSOCKET_CLIENT_ENTITY_CLASS[
+                        entity_info.__class__
+                    ](entity_info, self)
+                )
 
     def emit_platform_entity_event(self, event: EntityStateChangedEvent) -> None:
         """Proxy the firing of an entity event."""
@@ -1322,5 +1321,4 @@ def emit_platform_entity_event(self, event: EntityStateChangedEvent) -> None:
                 f"Entity not found: {event.platform}.{event.unique_id}",
             )
         entity.state = event.state
-        entity.maybe_emit_state_changed_event()
         self.emit(f"{event.unique_id}_{event.event}", event)