Compute primary entities (#298)

Author: puddly

tests/test_device.py

@@ -21,6 +21,7 @@
     SIG_EP_OUTPUT,
     SIG_EP_TYPE,
     create_mock_zigpy_device,
+    get_entity,
     join_zigpy_device,
     zigpy_device_from_json,
 )
@@ -34,6 +35,9 @@
     UNKNOWN,
 )
 from zha.application.gateway import Gateway
+from zha.application.platforms import PlatformEntity
+from zha.application.platforms.binary_sensor import IASZone
+from zha.application.platforms.light import Light
 from zha.application.platforms.sensor import LQISensor, RSSISensor
 from zha.application.platforms.switch import Switch
 from zha.exceptions import ZHAException
@@ -853,3 +857,63 @@ async def test_quirks_v2_device_alerts(zha_gateway: Gateway) -> None:
     assert zha_device.device_alerts == (
         DeviceAlertMetadata(level=DeviceAlertLevel.WARNING, message="Test warning"),
     )
+
+
+@pytest.mark.parametrize(
+    ("json_path", "primary_platform", "primary_entity_type"),
+    [
+        # Light bulb
+        (
+            "tests/data/devices/ikea-of-sweden-tradfri-bulb-gu10-ws-400lm.json",
+            Platform.LIGHT,
+            Light,
+        ),
+        # Night light with a bulb and a motion sensor
+        (
+            "tests/data/devices/third-reality-inc-3rsnl02043z.json",
+            Platform.LIGHT,
+            Light,
+        ),
+        # Door sensor
+        (
+            "tests/data/devices/centralite-3320-l.json",
+            Platform.BINARY_SENSOR,
+            IASZone,
+        ),
+        # Smart plug with energy monitoring
+        (
+            "tests/data/devices/innr-sp-234.json",
+            Platform.SWITCH,
+            Switch,
+        ),
+        # Atmosphere sensor with humidity, temperature, and pressure
+        (
+            "tests/data/devices/lumi-lumi-weather.json",
+            None,
+            None,
+        ),
+    ],
+)
+async def test_primary_entity_computation(
+    json_path: str,
+    primary_platform: Platform | None,
+    primary_entity_type: PlatformEntity | None,
+    zha_gateway: Gateway,
+) -> None:
+    """Test primary entity computation."""
+
+    zigpy_dev = await zigpy_device_from_json(
+        zha_gateway.application_controller,
+        json_path,
+    )
+    zha_device = await join_zigpy_device(zha_gateway, zigpy_dev)
+
+    # There is a single light entity
+    primary = [e for e in zha_device.platform_entities.values() if e.primary]
+
+    if primary_platform is None:
+        assert not primary
+    else:
+        assert primary == [
+            get_entity(zha_device, primary_platform, entity_type=primary_entity_type)
+        ]

zha/application/platforms/__init__.py

@@ -59,6 +59,7 @@ class BaseEntityInfo:
     entity_category: str | None
     entity_registry_enabled_default: bool
     enabled: bool = True
+    primary: bool
 
     # For platform entities
     cluster_handlers: list[ClusterHandlerInfo]
@@ -119,6 +120,11 @@ class BaseEntity(LogMixin, EventBase):
     _attr_state_class: str | None
     _attr_enabled: bool = True
     _attr_always_supported: bool = False
+    _attr_primary: bool = False
+
+    # When two entities both want to be primary, the one with the higher weight will be
+    # chosen. If there is a tie, both lose.
+    _attr_primary_weight: int = 0
 
     def __init__(self, unique_id: str) -> None:
         """Initialize the platform entity."""
@@ -156,6 +162,21 @@ def enabled(self, value: bool) -> None:
         """Set the entity enabled state."""
         self._attr_enabled = value
 
+    @property
+    def primary(self) -> bool:
+        """Return if the entity is the primary device control."""
+        return self._attr_primary
+
+    @primary.setter
+    def primary(self, value: bool) -> None:
+        """Set the entity as the primary device control."""
+        self._attr_primary = value
+
+    @property
+    def primary_weight(self) -> int:
+        """Return the primary weight of the entity."""
+        return self._attr_primary_weight
+
     @property
     def fallback_name(self) -> str | None:
         """Return the entity fallback name for when a translation key is unavailable."""
@@ -230,6 +251,7 @@ def info_object(self) -> BaseEntityInfo:
             entity_category=self.entity_category,
             entity_registry_enabled_default=self.entity_registry_enabled_default,
             enabled=self.enabled,
+            primary=self.primary,
             # Set by platform entities
             cluster_handlers=[],
             device_ieee=None,

zha/application/platforms/binary_sensor/__init__.py

@@ -167,13 +167,15 @@ class Occupancy(BinarySensor):
 
     _attribute_name = "occupancy"
     _attr_device_class: BinarySensorDeviceClass = BinarySensorDeviceClass.OCCUPANCY
+    _attr_primary_weight = 2
 
 
 @MULTI_MATCH(cluster_handler_names=CLUSTER_HANDLER_HUE_OCCUPANCY)
 class HueOccupancy(Occupancy):
     """ZHA Hue occupancy."""
 
     _attr_device_class: BinarySensorDeviceClass = BinarySensorDeviceClass.OCCUPANCY
+    _attr_primary_weight = 3
 
 
 @STRICT_MATCH(cluster_handler_names=CLUSTER_HANDLER_ON_OFF)
@@ -182,6 +184,7 @@ class Opening(BinarySensor):
 
     _attribute_name = "on_off"
     _attr_device_class: BinarySensorDeviceClass = BinarySensorDeviceClass.OPENING
+    _attr_primary_weight = 1
 
 
 @MULTI_MATCH(cluster_handler_names=CLUSTER_HANDLER_BINARY_INPUT)
@@ -215,6 +218,7 @@ class IASZone(BinarySensor):
     """ZHA IAS BinarySensor."""
 
     _attribute_name = "zone_status"
+    _attr_primary_weight = 3
 
     def __init__(
         self,
@@ -259,6 +263,7 @@ class SinopeLeakStatus(BinarySensor):
 
     _attribute_name = "leak_status"
     _attr_device_class = BinarySensorDeviceClass.MOISTURE
+    _attr_primary_weight = 1
 
 
 @MULTI_MATCH(

zha/application/platforms/climate/__init__.py

@@ -94,6 +94,7 @@ class Thermostat(PlatformEntity):
         ATTR_UNOCCP_COOL_SETPT,
         ATTR_UNOCCP_HEAT_SETPT,
     }
+    _attr_primary_weight = 10
 
     def __init__(
         self,

zha/application/platforms/cover/__init__.py

@@ -61,6 +61,7 @@ class Cover(PlatformEntity):
         "target_lift_position",
         "target_tilt_position",
     }
+    _attr_primary_weight = 10
 
     def __init__(
         self,
@@ -379,6 +380,7 @@ class Shade(PlatformEntity):
 
     _attr_device_class = CoverDeviceClass.SHADE
     _attr_translation_key: str = "shade"
+    _attr_primary_weight = 10
 
     def __init__(
         self,

zha/application/platforms/fan/__init__.py

@@ -80,6 +80,7 @@ class BaseFan(BaseEntity):
         | FanEntityFeature.TURN_ON
     )
     _attr_translation_key: str = "fan"
+    _attr_primary_weight = 10
 
     @functools.cached_property
     def preset_modes(self) -> list[str]:

zha/application/platforms/light/__init__.py

@@ -108,6 +108,7 @@ class BaseLight(BaseEntity, ABC):
         "off_with_transition",
         "off_brightness",
     }
+    _attr_primary_weight = 10
 
     def __init__(self, *args, **kwargs):
         """Initialize the light."""

zha/application/platforms/lock/__init__.py

@@ -36,6 +36,7 @@ class DoorLock(PlatformEntity):
 
     PLATFORM = Platform.LOCK
     _attr_translation_key: str = "door_lock"
+    _attr_primary_weight = 10
 
     def __init__(
         self,

zha/application/platforms/sensor/__init__.py

@@ -840,6 +840,7 @@ class Humidity(Sensor):
     _attr_state_class: SensorStateClass = SensorStateClass.MEASUREMENT
     _divisor = 100
     _attr_native_unit_of_measurement = PERCENTAGE
+    _attr_primary_weight = 1
 
 
 @MULTI_MATCH(cluster_handler_names=CLUSTER_HANDLER_SOIL_MOISTURE)
@@ -852,6 +853,7 @@ class SoilMoisture(Sensor):
     _attr_translation_key: str = "soil_moisture"
     _divisor = 100
     _attr_native_unit_of_measurement = PERCENTAGE
+    _attr_primary_weight = 1
 
 
 @MULTI_MATCH(cluster_handler_names=CLUSTER_HANDLER_LEAF_WETNESS)
@@ -864,6 +866,7 @@ class LeafWetness(Sensor):
     _attr_translation_key: str = "leaf_wetness"
     _divisor = 100
     _attr_native_unit_of_measurement = PERCENTAGE
+    _attr_primary_weight = 1
 
 
 @MULTI_MATCH(cluster_handler_names=CLUSTER_HANDLER_ILLUMINANCE)
@@ -874,6 +877,7 @@ class Illuminance(Sensor):
     _attr_device_class: SensorDeviceClass = SensorDeviceClass.ILLUMINANCE
     _attr_state_class: SensorStateClass = SensorStateClass.MEASUREMENT
     _attr_native_unit_of_measurement = LIGHT_LUX
+    _attr_primary_weight = 1
 
     def formatter(self, value: int) -> int | None:
         """Convert illumination data."""
@@ -911,6 +915,7 @@ class SmartEnergyMetering(PollableSensor):
         "status",
         "zcl_unit_of_measurement",
     }
+    _attr_primary_weight = 1
 
     _ENTITY_DESCRIPTION_MAP = {
         0x00: SmartEnergyMeteringEntityDescription(
@@ -1221,6 +1226,7 @@ class Pressure(Sensor):
     _attr_state_class: SensorStateClass = SensorStateClass.MEASUREMENT
     _decimals = 0
     _attr_native_unit_of_measurement = UnitOfPressure.HPA
+    _attr_primary_weight = 1
 
 
 @MULTI_MATCH(cluster_handler_names=CLUSTER_HANDLER_FLOW)
@@ -1232,6 +1238,7 @@ class Flow(Sensor):
     _attr_state_class: SensorStateClass = SensorStateClass.MEASUREMENT
     _divisor = 10
     _attr_native_unit_of_measurement = UnitOfVolumeFlowRate.CUBIC_METERS_PER_HOUR
+    _attr_primary_weight = 1
 
     def formatter(self, value: int) -> datetime | int | float | str | None:
         """Handle unknown value state."""
@@ -1249,6 +1256,7 @@ class Temperature(Sensor):
     _attr_state_class: SensorStateClass = SensorStateClass.MEASUREMENT
     _divisor = 100
     _attr_native_unit_of_measurement = UnitOfTemperature.CELSIUS
+    _attr_primary_weight = 1
 
 
 @MULTI_MATCH(cluster_handler_names=CLUSTER_HANDLER_DEVICE_TEMPERATURE)
@@ -1262,6 +1270,7 @@ class DeviceTemperature(Sensor):
     _divisor = 100
     _attr_native_unit_of_measurement = UnitOfTemperature.CELSIUS
     _attr_entity_category = EntityCategory.DIAGNOSTIC
+    _attr_primary_weight = 1
 
 
 @MULTI_MATCH(cluster_handler_names=CLUSTER_HANDLER_INOVELLI)
@@ -1304,6 +1313,7 @@ class CarbonDioxideConcentration(Sensor):
     _decimals = 0
     _multiplier = 1e6
     _attr_native_unit_of_measurement = CONCENTRATION_PARTS_PER_MILLION
+    _attr_primary_weight = 1
 
 
 @MULTI_MATCH(cluster_handler_names="carbon_monoxide_concentration")
@@ -1316,6 +1326,7 @@ class CarbonMonoxideConcentration(Sensor):
     _decimals = 0
     _multiplier = 1e6
     _attr_native_unit_of_measurement = CONCENTRATION_PARTS_PER_MILLION
+    _attr_primary_weight = 1
 
 
 @MULTI_MATCH(generic_ids="cluster_handler_0x042e", stop_on_match_group="voc_level")
@@ -1329,6 +1340,7 @@ class VOCLevel(Sensor):
     _decimals = 0
     _multiplier = 1e6
     _attr_native_unit_of_measurement = CONCENTRATION_MICROGRAMS_PER_CUBIC_METER
+    _attr_primary_weight = 1
 
 
 @MULTI_MATCH(
@@ -1347,6 +1359,7 @@ class PPBVOCLevel(Sensor):
     _decimals = 0
     _multiplier = 1
     _attr_native_unit_of_measurement = CONCENTRATION_PARTS_PER_BILLION
+    _attr_primary_weight = 1
 
 
 @MULTI_MATCH(cluster_handler_names="pm25")
@@ -1359,6 +1372,7 @@ class PM25(Sensor):
     _decimals = 0
     _multiplier = 1
     _attr_native_unit_of_measurement = CONCENTRATION_MICROGRAMS_PER_CUBIC_METER
+    _attr_primary_weight = 1
 
 
 @MULTI_MATCH(cluster_handler_names="formaldehyde_concentration")
@@ -1371,6 +1385,7 @@ class FormaldehydeConcentration(Sensor):
     _decimals = 0
     _multiplier = 1e6
     _attr_native_unit_of_measurement = CONCENTRATION_PARTS_PER_MILLION
+    _attr_primary_weight = 1
 
 
 @MULTI_MATCH(

zha/application/platforms/siren.py

@@ -68,6 +68,7 @@ class Siren(PlatformEntity):
 
     PLATFORM = Platform.SIREN
     _attr_fallback_name: str = "Siren"
+    _attr_primary_weight = 10
 
     def __init__(
         self,