Separate computed primary entity state from _attr_primary

tests/test_device.py

@@ -18,7 +18,7 @@
 import zigpy.types
 from zigpy.typing import UNDEFINED
 from zigpy.zcl import ClusterType
-from zigpy.zcl.clusters import general
+from zigpy.zcl.clusters import general, security
 from zigpy.zcl.clusters.general import Ota, PowerConfiguration
 from zigpy.zcl.clusters.lighting import Color
 from zigpy.zcl.clusters.measurement import CarbonDioxideConcentration
@@ -973,6 +973,122 @@ async def test_primary_entity_weight_0_not_elected(zha_gateway: Gateway) -> None
     assert not battery.primary
 
 
+async def test_primary_entity_reelection(zha_gateway: Gateway) -> None:
+    """Test election losers are not permanently excluded from later elections."""
+
+    # A smart plug with an IAS zone
+    zigpy_dev = create_mock_zigpy_device(
+        zha_gateway,
+        {
+            1: {
+                SIG_EP_INPUT: [
+                    general.OnOff.cluster_id,
+                    security.IasZone.cluster_id,
+                ],
+                SIG_EP_OUTPUT: [],
+                SIG_EP_TYPE: zigpy.profiles.zha.DeviceType.SMART_PLUG,
+                SIG_EP_PROFILE: zigpy.profiles.zha.PROFILE_ID,
+            }
+        },
+    )
+    zha_device = await join_zigpy_device(zha_gateway, zigpy_dev)
+
+    switch = get_entity(zha_device, Platform.SWITCH, entity_type=Switch)
+    ias_zone = get_entity(zha_device, Platform.BINARY_SENSOR, entity_type=IASZone)
+
+    assert switch.primary
+    assert not ias_zone.primary
+
+    # When the `on_off` attribute becomes unsupported, the switch is removed and a
+    # re-election elects the runner-up instead of permanently leaving the device
+    # without a primary entity
+    zigpy_dev.endpoints[1].on_off.add_unsupported_attribute(
+        general.OnOff.AttributeDefs.on_off.id
+    )
+    await zha_device.recompute_entities()
+
+    assert (Platform.SWITCH, switch.unique_id) not in zha_device.platform_entities
+    assert ias_zone.primary
+
+    # Writing a value clears the unsupported flag; the rediscovered switch wins
+    # back the election
+    zigpy_dev.endpoints[1].on_off.update_attribute(
+        general.OnOff.AttributeDefs.on_off.id, zigpy.types.Bool.false
+    )
+    await zha_device.recompute_entities()
+
+    switch = get_entity(zha_device, Platform.SWITCH, entity_type=Switch)
+    assert switch.primary
+    assert not ias_zone.primary
+
+
+async def test_primary_entity_election_disabled_winner(zha_gateway: Gateway) -> None:
+    """Test a disabled previous winner does not keep stale computed primary state."""
+
+    # A smart plug with an IAS zone
+    zigpy_dev = create_mock_zigpy_device(
+        zha_gateway,
+        {
+            1: {
+                SIG_EP_INPUT: [
+                    general.OnOff.cluster_id,
+                    security.IasZone.cluster_id,
+                ],
+                SIG_EP_OUTPUT: [],
+                SIG_EP_TYPE: zigpy.profiles.zha.DeviceType.SMART_PLUG,
+                SIG_EP_PROFILE: zigpy.profiles.zha.PROFILE_ID,
+            }
+        },
+    )
+    zha_device = await join_zigpy_device(zha_gateway, zigpy_dev)
+
+    switch = get_entity(zha_device, Platform.SWITCH, entity_type=Switch)
+    ias_zone = get_entity(zha_device, Platform.BINARY_SENSOR, entity_type=IASZone)
+
+    assert switch.primary
+    assert not ias_zone.primary
+
+    # When the switch is disabled, it is no longer an election candidate and must
+    # not keep its computed primary state, so only the runner-up is primary
+    switch.disable()
+    await zha_device.recompute_entities()
+
+    assert not switch.primary
+    assert ias_zone.primary
+
+    # When the switch is re-enabled, it wins back the election
+    switch.enable()
+    await zha_device.recompute_entities()
+
+    assert switch.primary
+    assert not ias_zone.primary
+
+
+async def test_primary_entity_election_explicit_primary_takes_over(
+    zha_gateway: Gateway,
+) -> None:
+    """Test an explicitly primary entity replaces a previously computed winner."""
+
+    # Night light with a bulb and a motion sensor
+    zigpy_dev = await zigpy_device_from_json(
+        zha_gateway.application_controller,
+        "tests/data/devices/third-reality-inc-3rsnl02043z-0x0000003c.json",
+    )
+    zha_device = await join_zigpy_device(zha_gateway, zigpy_dev)
+
+    light = get_entity(zha_device, Platform.LIGHT, entity_type=Light)
+    motion = get_entity(zha_device, Platform.BINARY_SENSOR, entity_type=IASZone)
+    assert light.primary
+
+    # Mark the motion sensor as explicitly primary, like a quirk would.
+    # Recomputing the entities re-runs the primary entity election.
+    motion._attr_primary = True
+    await zha_device.recompute_entities()
+
+    assert motion.primary
+    assert not light.primary
+
+
 async def test_quirks_v2_primary_entity(zha_gateway: Gateway) -> None:
     """Test quirks v2 primary entity."""
     registry = DeviceRegistry()

zha/application/platforms/__init__.py

@@ -279,6 +279,10 @@ async def async_initialize_cluster(self, cluster: Any) -> None:
     _attr_enabled: bool = True
     _attr_extra_state_attribute_names: set[str] | None = None
     _attr_always_supported: bool = False
+
+    # Explicitly marks the entity as (not) primary, set by entity classes and quirks.
+    # It takes precedence over (and is never overwritten by) the weight-based primary
+    # entity election, whose result is stored separately.
     _attr_primary: bool | None = None
 
     # When two entities both want to be primary, the one with the higher weight will be
@@ -291,6 +295,7 @@ def __init__(self, unique_id: str) -> None:
 
         self._unique_id: str = unique_id
         self._migrate_unique_ids: list[str] = []
+        self.__computed_primary: bool = False
 
         self.__previous_state: Any = None
         self._tracked_tasks: list[asyncio.Task] = []
@@ -329,15 +334,19 @@ def enabled(self, value: bool) -> None:
     @property
     def primary(self) -> bool:
         """Return if the entity is the primary device control."""
-        if self._attr_primary is None:
-            return False
+        if self._attr_primary is not None:
+            return self._attr_primary
 
-        return self._attr_primary
+        return self.__computed_primary
 
     @primary.setter
-    def primary(self, value: bool | None) -> None:
-        """Set the entity as the primary device control."""
-        self._attr_primary = value
+    def primary(self, value: bool) -> None:
+        """Set the computed primary state from the primary entity election.
+
+        An explicit `_attr_primary` always takes precedence and is never
+        overwritten by the election.
+        """
+        self.__computed_primary = value
 
     @property
     def primary_weight(self) -> int:

zha/zigbee/device.py

@@ -1612,8 +1612,14 @@ def log(self, level: int, msg: str, *args: Any, **kwargs: Any) -> None:
     def _compute_primary_entity(self, entities: Sequence[PlatformEntity]) -> None:
         """Compute the primary entity from a given set of entities."""
 
+        # Clear all previously computed primary state up front, so a stale winner
+        # cannot survive a re-election on any code path below (e.g. after it was
+        # disabled and thus is no longer a candidate)
+        for entity in entities:
+            entity.primary = False
+
         # First, check if any entity is explicitly primary
-        explicitly_primary = [entity for entity in entities if entity.primary]
+        explicitly_primary = [entity for entity in entities if entity._attr_primary]
 
         if len(explicitly_primary) == 1:
             self.debug(
@@ -1643,19 +1649,12 @@ def _compute_primary_entity(self, entities: Sequence[PlatformEntity]) -> None:
         # We have a clear winner
         if not others or winner.primary_weight > others[0].primary_weight:
             winner.primary = True
-
-            for entity in others:
-                entity.primary = False
-
             return
 
         self.debug(
             "Primary entity tie between %s and %s, no primary entity", winner, others[0]
         )
 
-        for entity in candidates:
-            entity.primary = False
-
     def get_diagnostics_json(self):
         """Get ZHA device information."""