From e42993a2c0f76cb0e9e31682249a1e2242d69858 Mon Sep 17 00:00:00 2001
From: puddly <32534428+puddly@users.noreply.github.com>
Date: Wed, 27 Nov 2024 15:12:53 -0500
Subject: [PATCH] Add an explicit unit test

---
 tests/test_device.py | 64 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 64 insertions(+)

diff --git a/tests/test_device.py b/tests/test_device.py
index ef52b3e85..7b3af30af 100644
--- a/tests/test_device.py
+++ b/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
@@ -820,3 +824,63 @@ async def test_quirks_v2_device_renaming(zha_gateway: Gateway) -> None:
     zha_device = await join_zigpy_device(zha_gateway, zigpy_dev)
     assert zha_device.model == "IRIS Keypad V2"
     assert zha_device.manufacturer == "Lowe's"
+
+
+@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)
+        ]