Make update_attribute manufacturer code aware (zigpy#1760)

* Make `update_attribute` manufacturer code aware

* Add test for manufacturer code `update_attribute` handling

* Remove legacy cache fallback in `_legacy_apply_quirk_attribute_update`

* Adjust comments

* Use `mock_attribute_report` in test

* Mark attribute as unsupported initially

Author: TheJulianJES

tests/test_zcl.py

@@ -2490,3 +2490,101 @@ class ServerCommandDefs(Basic.ServerCommandDefs):
         (TestCluster.ServerCommandDefs.reset_fact_default, None),
     ]:
         assert cluster._get_effective_manufacturer_code(definition) is expected
+
+
+async def test_quirk_manufacturer_code_context_isolation(app_mock) -> None:
+    """Test that manufacturer code context is properly handled in _update_attribute.
+
+    When a manufacturer-specific attribute is reported and the cluster has multiple
+    attributes sharing the same ID (with different manufacturer codes), the
+    _update_attribute call must use the correct manufacturer code. This tests that:
+    1. The value is stored directly in the typed cache (not via legacy cache fallback)
+    2. Other attributes updated by quirks don't inherit the manufacturer code context
+    """
+
+    class TestCluster(zcl.Cluster):
+        cluster_id = 0xABCD
+        ep_attribute = "test_cluster"
+        _skip_registry = True
+
+        class AttributeDefs(zcl.foundation.BaseAttributeDefs):
+            # Two attributes sharing the same ID with different manufacturer codes
+            manuf_attr = foundation.ZCLAttributeDef(
+                id=0x0001,
+                type=t.uint8_t,
+                manufacturer_code=0x1234,
+            )
+            standard_attr = foundation.ZCLAttributeDef(
+                id=0x0001,
+                type=t.uint8_t,
+                manufacturer_code=None,
+            )
+            # A different attribute that the quirk will also update
+            other_attr = foundation.ZCLAttributeDef(
+                id=0x0002,
+                type=t.uint8_t,
+            )
+
+        def _update_attribute(self, attrid, value):
+            super()._update_attribute(attrid, value)
+
+            # When updating the manufacturer-specific attribute, also update other_attr
+            if attrid == self.AttributeDefs.manuf_attr.id:
+                super()._update_attribute(self.AttributeDefs.other_attr.id, 99)
+
+    dev = add_initialized_device(app_mock, nwk=0x1234, ieee=make_ieee(1))
+    cluster = TestCluster(dev.endpoints[1])
+    dev.endpoints[1].add_input_cluster(TestCluster.cluster_id, cluster)
+
+    events = []
+    cluster.on_event(AttributeReportedEvent.event_type, events.append)
+    cluster.on_event(AttributeUpdatedEvent.event_type, events.append)
+
+    # The attribute is currently marked as unsupported
+    cluster.add_unsupported_attribute(TestCluster.AttributeDefs.manuf_attr)
+
+    # Report the manufacturer-specific attribute
+    await mock_attribute_report(
+        cluster, {TestCluster.AttributeDefs.manuf_attr: t.uint8_t(42)}
+    )
+
+    # The legacy cache should not contain the attribute, as the typed cache was used
+    assert 0x0001 not in cluster._attr_cache._legacy_cache
+
+    # Verify that the manufacturer-specific attribute was stored correctly
+    assert cluster._attr_cache.get_value(TestCluster.AttributeDefs.manuf_attr) == 42
+
+    # Verify that the standard attribute (same ID, no manufacturer code) was NOT updated
+    with pytest.raises(KeyError):
+        cluster._attr_cache.get_value(TestCluster.AttributeDefs.standard_attr)
+
+    # Verify that other_attr was updated (by the quirk) without manufacturer code context
+    assert cluster._attr_cache.get_value(TestCluster.AttributeDefs.other_attr) == 99
+
+    # Verify the events have the correct manufacturer codes
+    assert len(events) == 2
+
+    # First event: other_attr updated by quirk (should have no manufacturer code)
+    assert events[0] == AttributeUpdatedEvent(
+        device_ieee=str(dev.ieee),
+        endpoint_id=1,
+        cluster_type=zcl.ClusterType.Server,
+        cluster_id=TestCluster.cluster_id,
+        attribute_name="other_attr",
+        attribute_id=TestCluster.AttributeDefs.other_attr.id,
+        manufacturer_code=None,
+        value=99,
+    )
+
+    # Second event: manuf_attr reported (should have the manufacturer code)
+    assert events[1] == AttributeReportedEvent(
+        device_ieee=str(dev.ieee),
+        endpoint_id=1,
+        cluster_type=zcl.ClusterType.Server,
+        cluster_id=TestCluster.cluster_id,
+        attribute_name="manuf_attr",
+        attribute_id=TestCluster.AttributeDefs.manuf_attr.id,
+        manufacturer_code=0x1234,
+        raw_value=42,
+        value=42,
+    )

zigpy/zcl/__init__.py

@@ -38,6 +38,14 @@
     "_suppressed_attribute_updates", default=frozenset()
 )
 
+# Tracks the (attribute_id, manufacturer_code) for the current attribute update operation.
+# Used to preserve manufacturer code context when calling _update_attribute,
+# so that manufacturer-specific attributes with conflicting IDs are stored correctly.
+# The manufacturer code is only applied when the attribute ID matches the original.
+_attribute_update_context: ContextVar[tuple[int, int | None] | None] = ContextVar(
+    "_attribute_update_context", default=None
+)
+
 
 @contextlib.contextmanager
 def _suppress_attribute_update_event(
@@ -53,6 +61,25 @@ def _suppress_attribute_update_event(
         _suppressed_attribute_updates.reset(token)
 
 
+@contextlib.contextmanager
+def _set_attribute_update_context(
+    attrid: int,
+    manufacturer_code: int | None,
+) -> Generator[None, None, None]:
+    """Set the attribute update context for preserving manufacturer code.
+
+    This allows quirks that call _update_attribute to preserve the manufacturer code,
+    ensuring manufacturer-specific attributes are stored with the correct cache key.
+    The manufacturer code is only applied when the attribute ID matches.
+    """
+    token = _attribute_update_context.set((attrid, manufacturer_code))
+
+    try:
+        yield
+    finally:
+        _attribute_update_context.reset(token)
+
+
 class ClusterType(enum.IntEnum):
     Server = 0
     Client = 1
@@ -443,23 +470,19 @@ def _legacy_apply_quirk_attribute_update(
 
         Returns None if the quirk swallowed the attribute (no super() call).
         """
-        with _suppress_attribute_update_event(self.cluster_id, attr_def.id):
+        manufacturer_code = self._get_effective_manufacturer_code(attr_def)
+
+        with (
+            _suppress_attribute_update_event(self.cluster_id, attr_def.id),
+            _set_attribute_update_context(attr_def.id, manufacturer_code),
+        ):
             self._update_attribute(attr_def.id, value)
 
         try:
             return self._attr_cache.get_value(attr_def)
         except KeyError:
-            pass
-
-        # When multiple attrs share an ID (different manufacturer codes),
-        # `_update_attribute` stores in legacy cache. Move it to typed cache.
-        if attr_def.id in self._attr_cache._legacy_cache:
-            cached_value = self._attr_cache._legacy_cache.pop(attr_def.id).value
-            self._attr_cache.set_value(attr_def, cached_value)
-            return cached_value
-
-        # Quirk swallowed the attribute
-        return None
+            # Quirk swallowed the attribute
+            return None
 
     @classmethod
     def find_attributes(
@@ -1138,8 +1161,17 @@ def _update_attribute(
         # other clusters or attributes to emit their own events.
         suppressed = (self.cluster_id, attrid) in _suppressed_attribute_updates.get()
 
+        # Get the manufacturer code from context if set (set by
+        # _legacy_apply_quirk_attribute_update to preserve manufacturer code).
+        # Only apply when the attribute ID matches the original to avoid affecting
+        # other attributes that quirks may update.
+        ctx = _attribute_update_context.get()
+
         try:
-            attr_def = self.find_attribute(attrid)
+            if ctx is not None and ctx[0] == attrid:
+                attr_def = self.find_attribute(attrid, manufacturer_code=ctx[1])
+            else:
+                attr_def = self.find_attribute(attrid)
         except KeyError:
             if value is not None:
                 self._attr_cache.set_legacy_value(attrid, value)