Drop strict mode by fixing violations

zigpy · Jan 30, 2025 · faa5e63 · faa5e63
1 parent fb3dfd6
commit faa5e63
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 9 deletions.
diff --git a/tests/common.py b/tests/common.py
@@ -229,7 +229,6 @@ def get_entity(
     exact_entity_type: type[BaseEntity] | None = None,
     qualifier: str | None = None,
     qualifier_func: Callable[[BaseEntity], bool] = lambda e: True,
-    strict: bool = False,
 ) -> PlatformEntity:
     """Get the first entity of the specified platform on the given device."""
     results = []
@@ -257,7 +256,7 @@ def get_entity(
             f"No {entity_type} entity found for platform {platform!r} on device {device}: {device.platform_entities}"
         )
 
-    if strict and len(results) != 1:
+    if len(results) != 1:
         raise KeyError(
             f"Multiple {entity_type} entities found for platform {platform!r} on device {device}: {results}"
         )

diff --git a/tests/test_button.py b/tests/test_button.py
@@ -143,7 +143,6 @@ async def test_frost_unlock(
         zha_device,
         platform=Platform.BUTTON,
         entity_type=FrostLockResetButton,
-        strict=True,
     )
     assert isinstance(entity, FrostLockResetButton)
 
@@ -253,7 +252,9 @@ async def test_quirks_command_button(
     """Test ZHA button platform."""
     zha_device, cluster = await custom_button_device(zha_gateway)
     assert cluster is not None
-    entity: PlatformEntity = get_entity(zha_device, platform=Platform.BUTTON)
+    entity: PlatformEntity = get_entity(
+        zha_device, platform=Platform.BUTTON, entity_type=Button
+    )
 
     with patch(
         "zigpy.zcl.Cluster.request",

diff --git a/tests/test_select.py b/tests/test_select.py
@@ -186,7 +186,11 @@ async def test_on_off_select_attribute_report_v2(
     cluster = zigpy_device.endpoints[1].opple_cluster
     assert isinstance(zha_device.device, CustomDeviceV2)
 
-    entity = get_entity(zha_device, platform=Platform.SELECT)
+    entity = get_entity(
+        zha_device,
+        platform=Platform.SELECT,
+        qualifier_func=lambda e: e.info_object.unique_id.endswith("motion_sensitivity"),
+    )
 
     # test that the state is in default medium state
     assert entity.state["state"] == AqaraMotionSensitivities.Medium.name

diff --git a/tests/test_sensor.py b/tests/test_sensor.py
@@ -1103,7 +1103,11 @@ async def test_elec_measurement_sensor_type(
 
     zha_dev = await join_zigpy_device(zha_gateway, zigpy_dev)
 
-    entity = get_entity(zha_dev, platform=Platform.SENSOR)
+    entity = get_entity(
+        zha_dev,
+        platform=Platform.SENSOR,
+        entity_type=sensor.ElectricalMeasurementApparentPower,
+    )
     assert entity.state["measurement_type"] == expected_type
 
 
@@ -1450,8 +1454,16 @@ async def test_state_class(
 
     zha_device, cluster = await zigpy_device_aqara_sensor_v2_mock(zha_gateway)
     assert isinstance(zha_device.device, CustomDeviceV2)
-    power_entity = get_entity(zha_device, platform=Platform.SENSOR, qualifier="power")
-    energy_entity = get_entity(zha_device, platform=Platform.SENSOR, qualifier="energy")
+    power_entity = get_entity(
+        zha_device,
+        platform=Platform.SENSOR,
+        qualifier_func=lambda e: e.info_object.unique_id.endswith("power"),
+    )
+    energy_entity = get_entity(
+        zha_device,
+        platform=Platform.SENSOR,
+        qualifier_func=lambda e: e.info_object.unique_id.endswith("energy"),
+    )
     energy_delivered_entity = get_entity(
         zha_device, platform=Platform.SENSOR, qualifier="energy_delivered"
     )
@@ -1518,7 +1530,13 @@ async def test_device_counter_sensors(zha_gateway: Gateway) -> None:
 
     coordinator = zha_gateway.coordinator_zha_device
     assert coordinator.is_coordinator
-    entity = get_entity(coordinator, platform=Platform.SENSOR)
+    entity = get_entity(
+        coordinator,
+        platform=Platform.SENSOR,
+        qualifier_func=lambda e: e.info_object.unique_id.endswith(
+            "ezsp_counters_counter_1"
+        ),
+    )
 
     assert entity.state["state"] == 1