[tests] Use the exact entity lookups in the integration tests (#21213)

tests/integration/state_utils.py

@@ -181,7 +181,9 @@ def require_entity[T: EntityInfo](
             else f"entity with '{object_id_substring}' in object_id"
         )
         type_info = f" of type {entity_type.__name__}" if entity_type else ""
-        raise AssertionError(f"{desc}{type_info} not found in entities")
+        raise AssertionError(
+            f"{desc}{type_info} not found; available: {[e.object_id for e in entities]}"
+        )
     return entity
 
 

tests/integration/test_api_message_size_batching.py

@@ -7,7 +7,7 @@ import asyncio
 from aioesphomeapi import EntityState, NumberInfo, SelectInfo, TextInfo, TextSensorInfo
 import pytest
 
-from .state_utils import assert_device_name
+from .state_utils import assert_device_name, require_entity
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -59,33 +59,13 @@ async def test_api_message_size_batching(
             f"Expected exactly 1 text input entity, got {len(text_inputs)}"
         )
 
-        # Collect all select entity object_ids for error messages
-        select_ids = [s.object_id for s in selects]
-
         # Find our specific test entities
-        small_select = None
-        medium_select = None
-        large_select = None
-
-        for select in selects:
-            if select.object_id == "small_select":
-                small_select = select
-            elif select.object_id == "medium_select":
-                medium_select = select
-            elif (
-                select.object_id
-                == "large_select_with_many_options_to_create_larger_payload"
-            ):
-                large_select = select
-
-        assert small_select is not None, (
-            f"Could not find small_select entity. Found: {select_ids}"
-        )
-        assert medium_select is not None, (
-            f"Could not find medium_select entity. Found: {select_ids}"
-        )
-        assert large_select is not None, (
-            f"Could not find large_select entity. Found: {select_ids}"
+        small_select = require_entity(selects, "small_select", exact=True)
+        medium_select = require_entity(selects, "medium_select", exact=True)
+        large_select = require_entity(
+            selects,
+            "large_select_with_many_options_to_create_larger_payload",
+            exact=True,
         )
 
         # Verify the selects have the expected number of options
@@ -99,44 +79,15 @@ async def test_api_message_size_batching(
             f"Expected 50 options for large_select, got {len(large_select.options)}"
         )
 
-        # Collect all text sensor object_ids for error messages
-        text_sensor_ids = [t.object_id for t in text_sensors]
-
         # Verify text sensors with different value lengths
-        short_text_sensor = None
-        medium_text_sensor = None
-        long_text_sensor = None
-
-        for text_sensor in text_sensors:
-            if text_sensor.object_id == "short_text_sensor":
-                short_text_sensor = text_sensor
-            elif text_sensor.object_id == "medium_text_sensor":
-                medium_text_sensor = text_sensor
-            elif text_sensor.object_id == "long_text_sensor_with_very_long_value":
-                long_text_sensor = text_sensor
-
-        assert short_text_sensor is not None, (
-            f"Could not find short_text_sensor. Found: {text_sensor_ids}"
-        )
-        assert medium_text_sensor is not None, (
-            f"Could not find medium_text_sensor. Found: {text_sensor_ids}"
-        )
-        assert long_text_sensor is not None, (
-            f"Could not find long_text_sensor. Found: {text_sensor_ids}"
+        require_entity(text_sensors, "short_text_sensor", exact=True)
+        require_entity(text_sensors, "medium_text_sensor", exact=True)
+        require_entity(
+            text_sensors, "long_text_sensor_with_very_long_value", exact=True
         )
 
         # Check text input which can have a long max_length
-        text_input = None
-        text_input_ids = [t.object_id for t in text_inputs]
-
-        for ti in text_inputs:
-            if ti.object_id == "test_text_input":
-                text_input = ti
-                break
-
-        assert text_input is not None, (
-            f"Could not find test_text_input. Found: {text_input_ids}"
-        )
+        text_input = require_entity(text_inputs, "test_text_input", exact=True)
         assert text_input.max_length == 255, (
             f"Expected max_length 255, got {text_input.max_length}"
         )

tests/integration/test_batch_delay_zero_rapid_transitions.py

@@ -8,6 +8,7 @@ import time
 from aioesphomeapi import BinarySensorInfo, BinarySensorState, EntityState
 import pytest
 
+from .state_utils import require_only_entity
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -32,8 +33,7 @@ async def test_batch_delay_zero_rapid_transitions(
 
         # Wait for entity info
         entity_info, _ = await client.list_entities_services()
-        binary_sensors = [e for e in entity_info if isinstance(e, BinarySensorInfo)]
-        assert len(binary_sensors) == 1, "Expected 1 binary sensor"
+        require_only_entity(entity_info, BinarySensorInfo)
 
         # Collect states for 2 seconds
         await asyncio.sleep(2.1)

tests/integration/test_build_info.py

@@ -10,7 +10,7 @@ import time
 from aioesphomeapi import EntityState, TextSensorState
 import pytest
 
-from .state_utils import assert_device_name
+from .state_utils import assert_device_name, require_entity
 from .types import APIClientConnectedFactory, RunCompiledFunction
 from .wait_utils import wait_or_fail
 
@@ -38,19 +38,11 @@ async def test_build_info(
         entities, _ = await client.list_entities_services()
 
         # Find our text sensors by object_id
-        config_hash_entity = next(
-            (e for e in entities if e.object_id == "config_hash"), None
+        config_hash_entity = require_entity(entities, "config_hash", exact=True)
+        build_time_entity = require_entity(entities, "build_time", exact=True)
+        build_time_str_entity = require_entity(
+            entities, "build_time_string", exact=True
         )
-        build_time_entity = next(
-            (e for e in entities if e.object_id == "build_time"), None
-        )
-        build_time_str_entity = next(
-            (e for e in entities if e.object_id == "build_time_string"), None
-        )
-
-        assert config_hash_entity is not None, "Config Hash sensor not found"
-        assert build_time_entity is not None, "Build Time sensor not found"
-        assert build_time_str_entity is not None, "Build Time String sensor not found"
 
         # Wait for all three text sensors to have valid states
         loop = asyncio.get_running_loop()

tests/integration/test_climate_custom_modes.py

@@ -8,7 +8,7 @@ import aioesphomeapi
 from aioesphomeapi import ClimateInfo, ClimatePreset, EntityState
 import pytest
 
-from .state_utils import InitialStateHelper
+from .state_utils import InitialStateHelper, require_only_entity
 from .types import APIClientConnectedFactory, RunCompiledFunction
 from .wait_utils import wait_or_fail
 
@@ -49,10 +49,7 @@ async def test_climate_custom_fan_modes_and_presets(
         # Get entities and set up state synchronization
         entities, services = await client.list_entities_services()
         initial_state_helper = InitialStateHelper(entities)
-        climate_infos = [e for e in entities if isinstance(e, ClimateInfo)]
-        assert len(climate_infos) == 1, "Expected exactly 1 climate entity"
-
-        test_climate = climate_infos[0]
+        test_climate = require_only_entity(entities, ClimateInfo)
 
         # Subscribe with the wrapper that filters initial states
         client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))

tests/integration/test_large_message_batching.py

@@ -5,7 +5,7 @@ from __future__ import annotations
 from aioesphomeapi import SelectInfo
 import pytest
 
-from .state_utils import assert_device_name
+from .state_utils import assert_device_name, require_entity
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -25,13 +25,9 @@ async def test_large_message_batching(
         entity_info, services = await client.list_entities_services()
 
         # Find our large select entity
-        large_select = None
-        for entity in entity_info:
-            if isinstance(entity, SelectInfo) and entity.object_id == "large_select":
-                large_select = entity
-                break
-
-        assert large_select is not None, "Could not find large_select entity"
+        large_select = require_entity(
+            entity_info, "large_select", SelectInfo, exact=True
+        )
 
         # Verify the select has all its options
         # We created 100 options with long names

tests/integration/test_legacy_climate_compat.py

@@ -16,7 +16,7 @@ import aioesphomeapi
 from aioesphomeapi import ClimateInfo
 import pytest
 
-from .state_utils import InitialStateHelper
+from .state_utils import InitialStateHelper, require_only_entity
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -33,12 +33,7 @@ async def test_legacy_climate_compat(
         entities, _ = await client.list_entities_services()
         initial_state_helper = InitialStateHelper(entities)
 
-        climate_infos = [e for e in entities if isinstance(e, ClimateInfo)]
-        assert len(climate_infos) == 1, (
-            f"Expected 1 climate entity, got {len(climate_infos)}"
-        )
-
-        test_climate = climate_infos[0]
+        test_climate = require_only_entity(entities, ClimateInfo)
 
         # Verify custom fan modes set via deprecated ClimateTraits setter are exposed
         assert set(test_climate.supported_custom_fan_modes) == {

tests/integration/test_legacy_fan_compat.py

@@ -14,6 +14,7 @@ import asyncio
 from aioesphomeapi import FanInfo, FanState
 import pytest
 
+from .state_utils import require_only_entity
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -27,10 +28,7 @@ async def test_legacy_fan_compat(
     async with run_compiled(yaml_config), api_client_connected() as client:
         entities, _ = await client.list_entities_services()
 
-        fan_infos = [e for e in entities if isinstance(e, FanInfo)]
-        assert len(fan_infos) == 1, f"Expected 1 fan entity, got {len(fan_infos)}"
-
-        test_fan = fan_infos[0]
+        test_fan = require_only_entity(entities, FanInfo)
 
         # Verify preset modes set via deprecated FanTraits setter are exposed
         assert set(test_fan.supported_preset_modes) == {

tests/integration/test_light_automations.py

@@ -8,6 +8,7 @@ import asyncio
 
 import pytest
 
+from .state_utils import require_entity
 from .types import APIClientConnectedFactory, RunCompiledFunction
 from .wait_utils import wait_or_fail
 
@@ -59,7 +60,7 @@ async def test_light_automations(
     ):
         # Get entities
         entities = await client.list_entities_services()
-        light = next(e for e in entities[0] if e.object_id == "test_light")
+        light = require_entity(entities[0], "test_light", exact=True)
 
         # Start counting on_state events now
         counting_enabled = True

tests/integration/test_light_binary_effect_off_phase.py

@@ -26,7 +26,7 @@ from typing import Any
 from aioesphomeapi import EntityState, LightState
 import pytest
 
-from .state_utils import InitialStateHelper
+from .state_utils import InitialStateHelper, require_entity
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 OUTPUT_PATTERN = re.compile(r"BINARY_OUTPUT:(YES|NO)")
@@ -64,7 +64,7 @@ async def test_light_binary_effect_off_phase(
         api_client_connected() as client,
     ):
         entities, _ = await client.list_entities_services()
-        light = next(e for e in entities if e.object_id == "test_binary_light")
+        light = require_entity(entities, "test_binary_light", exact=True)
 
         state_futures: dict[int, asyncio.Future[LightState]] = {}
 
@@ -163,7 +163,7 @@ async def test_light_binary_zero_brightness_is_recoverable(
         api_client_connected() as client,
     ):
         entities, _ = await client.list_entities_services()
-        light = next(e for e in entities if e.object_id == "test_binary_light")
+        light = require_entity(entities, "test_binary_light", exact=True)
 
         state_futures: dict[int, asyncio.Future[LightState]] = {}
 

tests/integration/test_light_control_action.py

@@ -10,6 +10,7 @@ from typing import Any
 
 import pytest
 
+from .state_utils import require_entity
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -31,7 +32,7 @@ async def test_light_control_action(
 
         # Get entities
         entities = await client.list_entities_services()
-        light = next(e for e in entities[0] if e.object_id == "test_light")
+        light = require_entity(entities[0], "test_light", exact=True)
         buttons = {e.name: e for e in entities[0] if hasattr(e, "name")}
 
         async def wait_for_state(key: int, timeout: float = 5.0) -> Any:

tests/integration/test_light_effect_zero_brightness.py

@@ -20,7 +20,7 @@ from typing import Any
 from aioesphomeapi import EntityState, LightState
 import pytest
 
-from .state_utils import InitialStateHelper
+from .state_utils import InitialStateHelper, require_entity
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -44,7 +44,7 @@ async def test_light_effect_zero_brightness(
         api_client_connected() as client,
     ):
         entities, _ = await client.list_entities_services()
-        light = next(e for e in entities if e.object_id == "test_pulse_light")
+        light = require_entity(entities, "test_pulse_light", exact=True)
 
         state_futures: dict[int, asyncio.Future[LightState]] = {}
 

tests/integration/test_light_repeat_effect.py

@@ -17,7 +17,7 @@ from typing import Any
 from aioesphomeapi import EntityState, LightState
 import pytest
 
-from .state_utils import InitialStateHelper
+from .state_utils import InitialStateHelper, require_entity
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -45,7 +45,7 @@ async def test_light_repeat_effect(
         api_client_connected() as client,
     ):
         entities, _ = await client.list_entities_services()
-        light = next(e for e in entities if e.object_id == "test_rgb_light")
+        light = require_entity(entities, "test_rgb_light", exact=True)
 
         state_futures: dict[int, asyncio.Future[LightState]] = {}
 

tests/integration/test_light_resume_effect.py

@@ -13,7 +13,7 @@ from typing import Any
 from aioesphomeapi import ColorMode, EntityState, LightInfo, LightState
 import pytest
 
-from .state_utils import InitialStateHelper
+from .state_utils import InitialStateHelper, require_entity
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -29,10 +29,8 @@ async def test_light_resume_effect(
         api_client_connected() as client,
     ):
         entities, _ = await client.list_entities_services()
-        light_resume = next(e for e in entities if e.object_id == "test_light_resume")
-        light_no_resume = next(
-            e for e in entities if e.object_id == "test_light_no_resume"
-        )
+        light_resume = require_entity(entities, "test_light_resume", exact=True)
+        light_no_resume = require_entity(entities, "test_light_no_resume", exact=True)
 
         state_futures: dict[int, asyncio.Future[LightState]] = {}
 

tests/integration/test_lock_automations.py

@@ -7,6 +7,7 @@ import asyncio
 
 import pytest
 
+from .state_utils import require_entity
 from .types import APIClientConnectedFactory, RunCompiledFunction
 from .wait_utils import wait_or_fail
 
@@ -40,7 +41,7 @@ async def test_lock_automations(
 
         # Get entities
         entities = await client.list_entities_services()
-        lock = next(e for e in entities[0] if e.object_id == "test_lock")
+        lock = require_entity(entities[0], "test_lock", exact=True)
 
         # Test 1: Lock - should trigger on_lock
         client.lock_command(key=lock.key, command=LockCommand.LOCK)

tests/integration/test_scheduler_defer_cancel.py

@@ -5,7 +5,7 @@ import asyncio
 from aioesphomeapi import EntityState, Event, EventInfo, UserService
 import pytest
 
-from .state_utils import assert_device_name
+from .state_utils import assert_device_name, require_entity
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -27,18 +27,12 @@ async def test_scheduler_defer_cancel(
         )
 
         # Find our test entities
-        test_complete_entity: EventInfo | None = None
-        test_result_entity: EventInfo | None = None
-
-        for entity in entity_info:
-            if isinstance(entity, EventInfo):
-                if entity.object_id == "test_complete":
-                    test_complete_entity = entity
-                elif entity.object_id == "test_result":
-                    test_result_entity = entity
-
-        assert test_complete_entity is not None, "test_complete event not found"
-        assert test_result_entity is not None, "test_result event not found"
+        test_complete_entity = require_entity(
+            entity_info, "test_complete", EventInfo, exact=True
+        )
+        test_result_entity = require_entity(
+            entity_info, "test_result", EventInfo, exact=True
+        )
 
         # Find our test service
         test_defer_cancel_service: UserService | None = None

tests/integration/test_scheduler_defer_fifo_simple.py

@@ -5,7 +5,7 @@ import asyncio
 from aioesphomeapi import EntityState, Event, EventInfo, UserService
 import pytest
 
-from .state_utils import assert_device_name
+from .state_utils import assert_device_name, require_entity
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -27,18 +27,12 @@ async def test_scheduler_defer_fifo_simple(
         )
 
         # Find our test entities
-        test_complete_entity: EventInfo | None = None
-        test_result_entity: EventInfo | None = None
-
-        for entity in entity_info:
-            if isinstance(entity, EventInfo):
-                if entity.object_id == "test_complete":
-                    test_complete_entity = entity
-                elif entity.object_id == "test_result":
-                    test_result_entity = entity
-
-        assert test_complete_entity is not None, "test_complete event not found"
-        assert test_result_entity is not None, "test_result event not found"
+        test_complete_entity = require_entity(
+            entity_info, "test_complete", EventInfo, exact=True
+        )
+        test_result_entity = require_entity(
+            entity_info, "test_result", EventInfo, exact=True
+        )
 
         # Find our test services
         test_set_timeout_service: UserService | None = None

tests/integration/test_sensor_filters_ring_buffer.py

@@ -7,7 +7,7 @@ import asyncio
 from aioesphomeapi import EntityState, SensorState
 import pytest
 
-from .state_utils import InitialStateHelper, build_key_to_entity_mapping
+from .state_utils import InitialStateHelper, build_key_to_entity_mapping, require_entity
 from .types import APIClientConnectedFactory, RunCompiledFunction
 from .wait_utils import wait_or_fail
 
@@ -90,11 +90,7 @@ async def test_sensor_filters_ring_buffer(
             pytest.fail("Timeout waiting for initial states")
 
         # Find the publish button
-        publish_button = next(
-            (e for e in entities if "publish_values_button" in e.object_id.lower()),
-            None,
-        )
-        assert publish_button is not None, "Publish Values Button not found"
+        publish_button = require_entity(entities, "publish_values_button")
 
         # Press the button to publish test values
         client.button_command(publish_button.key)

tests/integration/test_sensor_filters_sliding_window.py

@@ -7,7 +7,7 @@ import asyncio
 from aioesphomeapi import EntityState, SensorState
 import pytest
 
-from .state_utils import InitialStateHelper, build_key_to_entity_mapping
+from .state_utils import InitialStateHelper, build_key_to_entity_mapping, require_entity
 from .types import APIClientConnectedFactory, RunCompiledFunction
 from .wait_utils import wait_or_fail
 
@@ -122,11 +122,7 @@ async def test_sensor_filters_sliding_window(
             pytest.fail("Timeout waiting for initial states")
 
         # Find the publish button
-        publish_button = next(
-            (e for e in entities if "publish_values_button" in e.object_id.lower()),
-            None,
-        )
-        assert publish_button is not None, "Publish Values Button not found"
+        publish_button = require_entity(entities, "publish_values_button")
 
         # Press the button to publish test values
         client.button_command(publish_button.key)
@@ -251,11 +247,7 @@ async def test_sensor_filters_nan_handling(
             pytest.fail("Timeout waiting for initial states")
 
         # Find the publish button
-        publish_button = next(
-            (e for e in entities if "publish_nan_values_button" in e.object_id.lower()),
-            None,
-        )
-        assert publish_button is not None, "Publish NaN Values Button not found"
+        publish_button = require_entity(entities, "publish_nan_values_button")
 
         # Press the button
         client.button_command(publish_button.key)
@@ -350,11 +342,7 @@ async def test_sensor_filters_ring_buffer_wraparound(
             pytest.fail("Timeout waiting for initial state")
 
         # Find the publish button
-        publish_button = next(
-            (e for e in entities if "publish_wraparound_button" in e.object_id.lower()),
-            None,
-        )
-        assert publish_button is not None, "Publish Wraparound Button not found"
+        publish_button = require_entity(entities, "publish_wraparound_button")
 
         # Press the button
         # Will publish: 10, 20, 30, 5, 25, 15, 40, 35, 20

tests/integration/test_sensor_raw_state.py

@@ -13,6 +13,7 @@ import re
 from aioesphomeapi import APIClient, EntityInfo
 import pytest
 
+from .state_utils import require_entity
 from .types import APIClientConnectedFactory, RunCompiledFunction
 from .wait_utils import wait_or_fail
 
@@ -27,10 +28,7 @@ async def _press_and_read(
     future: asyncio.Future[tuple[float, float]],
     label: str,
 ) -> tuple[float, float]:
-    button = next(
-        (e for e in entities if button_object_id in e.object_id.lower()), None
-    )
-    assert button is not None, f"{button_object_id} not found"
+    button = require_entity(entities, button_object_id)
     client.button_command(button.key)
     return await wait_or_fail(future, 5.0, f"Timeout waiting for {label} log message")
 

tests/integration/test_sensor_timeout_filter.py

@@ -7,7 +7,7 @@ import asyncio
 from aioesphomeapi import EntityState, SensorState
 import pytest
 
-from .state_utils import InitialStateHelper, build_key_to_entity_mapping
+from .state_utils import InitialStateHelper, build_key_to_entity_mapping, require_entity
 from .types import APIClientConnectedFactory, RunCompiledFunction
 from .wait_utils import wait_or_fail
 
@@ -98,12 +98,7 @@ async def test_sensor_timeout_filter(
         # Helper to find buttons by object_id substring
         def find_button(object_id_substring: str) -> int:
             """Find a button by object_id substring and return its key."""
-            button = next(
-                (e for e in entities if object_id_substring in e.object_id.lower()),
-                None,
-            )
-            assert button is not None, f"Button '{object_id_substring}' not found"
-            return button.key
+            return require_entity(entities, object_id_substring).key
 
         # Find all test buttons
         test1_button_key = find_button("test_timeout_last_button")

tests/integration/test_template_climate_basic.py

@@ -26,7 +26,12 @@ from aioesphomeapi import (
 import pytest
 
 from .host_prefs import clear_host_prefs
-from .state_utils import InitialStateHelper, require_entity, wait_for_state
+from .state_utils import (
+    InitialStateHelper,
+    require_entity,
+    require_only_entity,
+    wait_for_state,
+)
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 DEVICE_NAME = "tmpl-clim-basic"
@@ -61,9 +66,7 @@ async def test_template_climate_basic(
 
         entities, _ = await client.list_entities_services()
         initial_state_helper = InitialStateHelper(entities)
-        climate_infos = [e for e in entities if isinstance(e, ClimateInfo)]
-        assert len(climate_infos) == 1, "Expected exactly 1 climate entity"
-        test_climate = climate_infos[0]
+        test_climate = require_only_entity(entities, ClimateInfo)
 
         # Advertised capabilities come straight from the supported_*/custom_* config lists.
         assert ClimateMode.OFF in test_climate.supported_modes

tests/integration/test_template_climate_custom_modes.py

@@ -13,7 +13,12 @@ from aioesphomeapi import ButtonInfo, ClimateInfo
 import pytest
 
 from .host_prefs import clear_host_prefs
-from .state_utils import InitialStateHelper, require_entity, wait_for_state
+from .state_utils import (
+    InitialStateHelper,
+    require_entity,
+    require_only_entity,
+    wait_for_state,
+)
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 DEVICE_NAME = "tmpl-clim-custom"
@@ -48,9 +53,7 @@ async def test_template_climate_custom_modes(
 
         entities, _ = await client.list_entities_services()
         initial_state_helper = InitialStateHelper(entities)
-        climate_infos = [e for e in entities if isinstance(e, ClimateInfo)]
-        assert len(climate_infos) == 1, "Expected exactly 1 climate entity"
-        test_climate = climate_infos[0]
+        test_climate = require_only_entity(entities, ClimateInfo)
 
         assert set(test_climate.supported_custom_fan_modes) == {
             "turbo",

tests/integration/test_template_climate_nonoptimistic.py

@@ -21,7 +21,12 @@ from aioesphomeapi import (
 import pytest
 
 from .host_prefs import clear_host_prefs
-from .state_utils import InitialStateHelper, require_entity, wait_for_state
+from .state_utils import (
+    InitialStateHelper,
+    require_entity,
+    require_only_entity,
+    wait_for_state,
+)
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 DEVICE_NAME = "tmpl-clim-nonopt"
@@ -54,9 +59,7 @@ async def test_template_climate_nonoptimistic(
 
         entities, _ = await client.list_entities_services()
         initial_state_helper = InitialStateHelper(entities)
-        climate_infos = [e for e in entities if isinstance(e, ClimateInfo)]
-        assert len(climate_infos) == 1, "Expected exactly 1 climate entity"
-        test_climate = climate_infos[0]
+        test_climate = require_only_entity(entities, ClimateInfo)
 
         confirm_button = require_entity(
             entities, "simulate_device_confirmation", ButtonInfo

tests/integration/test_template_climate_on_control_ordering.py

@@ -14,7 +14,7 @@ from aioesphomeapi import ClimateInfo, ClimateMode
 import pytest
 
 from .host_prefs import clear_host_prefs
-from .state_utils import InitialStateHelper, wait_for_state
+from .state_utils import InitialStateHelper, require_only_entity, wait_for_state
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 DEVICE_NAME = "tmpl-clim-oc-order"
@@ -49,9 +49,7 @@ async def test_template_climate_on_control_ordering(
 
         entities, _ = await client.list_entities_services()
         initial_state_helper = InitialStateHelper(entities)
-        climate_infos = [e for e in entities if isinstance(e, ClimateInfo)]
-        assert len(climate_infos) == 1, "Expected exactly 1 climate entity"
-        test_climate = climate_infos[0]
+        test_climate = require_only_entity(entities, ClimateInfo)
 
         client.subscribe_states(
             initial_state_helper.on_state_wrapper(lambda state: None)

tests/integration/test_template_climate_publish_all_fields.py

@@ -24,7 +24,12 @@ from aioesphomeapi import (
 import pytest
 
 from .host_prefs import clear_host_prefs
-from .state_utils import InitialStateHelper, require_entity, wait_for_state
+from .state_utils import (
+    InitialStateHelper,
+    require_entity,
+    require_only_entity,
+    wait_for_state,
+)
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 DEVICE_NAME = "tmpl-clim-publish-all"
@@ -58,8 +63,7 @@ async def test_template_climate_publish_all_fields(
 
         entities, _ = await client.list_entities_services()
         initial_state_helper = InitialStateHelper(entities)
-        climate_infos = [e for e in entities if isinstance(e, ClimateInfo)]
-        assert len(climate_infos) == 1, "Expected exactly 1 climate entity"
+        require_only_entity(entities, ClimateInfo)
 
         publish_button = require_entity(entities, "publish_all", ButtonInfo)
 

tests/integration/test_template_climate_sensor_push.py

@@ -15,7 +15,12 @@ from aioesphomeapi import ButtonInfo, ClimateInfo
 import pytest
 
 from .host_prefs import clear_host_prefs
-from .state_utils import InitialStateHelper, require_entity, wait_for_state
+from .state_utils import (
+    InitialStateHelper,
+    require_entity,
+    require_only_entity,
+    wait_for_state,
+)
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 DEVICE_NAME = "tmpl-clim-sensor-push"
@@ -43,9 +48,7 @@ async def test_template_climate_sensor_push(
 
         entities, _ = await client.list_entities_services()
         initial_state_helper = InitialStateHelper(entities)
-        climate_infos = [e for e in entities if isinstance(e, ClimateInfo)]
-        assert len(climate_infos) == 1, "Expected exactly 1 climate entity"
-        test_climate = climate_infos[0]
+        test_climate = require_only_entity(entities, ClimateInfo)
 
         publish_temp = require_entity(entities, "publish_temperature", ButtonInfo)
         publish_temp_same = require_entity(

tests/integration/test_template_climate_set_actions.py

@@ -20,7 +20,12 @@ from aioesphomeapi import (
 import pytest
 
 from .host_prefs import clear_host_prefs
-from .state_utils import InitialStateHelper, require_entity, wait_for_state
+from .state_utils import (
+    InitialStateHelper,
+    require_entity,
+    require_only_entity,
+    wait_for_state,
+)
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 DEVICE_NAME = "tmpl-clim-set-act"
@@ -50,9 +55,7 @@ async def test_template_climate_set_actions(
     ):
         entities, _ = await client.list_entities_services()
         initial_state_helper = InitialStateHelper(entities)
-        climate_infos = [e for e in entities if isinstance(e, ClimateInfo)]
-        assert len(climate_infos) == 1, "Expected exactly 1 climate entity"
-        test_climate = climate_infos[0]
+        test_climate = require_only_entity(entities, ClimateInfo)
 
         report_button = require_entity(entities, "report_device_state", ButtonInfo)
         unsupported_button = require_entity(

tests/integration/test_template_climate_two_point_temperature.py

@@ -15,7 +15,12 @@ from aioesphomeapi import ButtonInfo, ClimateInfo, ClimateMode
 import pytest
 
 from .host_prefs import clear_host_prefs
-from .state_utils import InitialStateHelper, require_entity, wait_for_state
+from .state_utils import (
+    InitialStateHelper,
+    require_entity,
+    require_only_entity,
+    wait_for_state,
+)
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 DEVICE_NAME = "tmpl-clim-two-point"
@@ -50,10 +55,7 @@ async def test_template_climate_two_point_temperature(
 
         entities, _ = await client.list_entities_services()
         initial_state_helper = InitialStateHelper(entities)
-        climate_infos = [e for e in entities if isinstance(e, ClimateInfo)]
-        assert len(climate_infos) == 1, "Expected exactly 1 climate entity"
-
-        test_climate = climate_infos[0]
+        test_climate = require_only_entity(entities, ClimateInfo)
         assert test_climate.name == "Test Two-Point Heatpump"
         assert test_climate.supports_two_point_target_temperature
         assert test_climate.supports_target_humidity

tests/integration/test_text_sensor_raw_state.py

@@ -14,7 +14,7 @@ import re
 
 import pytest
 
-from .state_utils import assert_device_name
+from .state_utils import assert_device_name, require_entity
 from .types import APIClientConnectedFactory, RunCompiledFunction
 from .wait_utils import wait_or_fail
 
@@ -138,17 +138,9 @@ async def test_text_sensor_raw_state(
         entities, _ = await client.list_entities_services()
 
         # Find the test buttons
-        no_filter_button = next(
-            (e for e in entities if "test_no_filter_button" in e.object_id.lower()),
-            None,
-        )
-        assert no_filter_button is not None, "Test No Filter Button not found"
+        no_filter_button = require_entity(entities, "test_no_filter_button")
 
-        with_filter_button = next(
-            (e for e in entities if "test_with_filter_button" in e.object_id.lower()),
-            None,
-        )
-        assert with_filter_button is not None, "Test With Filter Button not found"
+        with_filter_button = require_entity(entities, "test_with_filter_button")
 
         # Test 1: Text sensor without filters
         # get_raw_state() should return the same as state
@@ -186,11 +178,7 @@ async def test_text_sensor_raw_state(
 
         # Test 3: Append filter (StringRef-based)
         # "test" + " suffix" = "test suffix"
-        append_button = next(
-            (e for e in entities if "test_append_button" in e.object_id.lower()),
-            None,
-        )
-        assert append_button is not None, "Test Append Button not found"
+        append_button = require_entity(entities, "test_append_button")
         client.button_command(append_button.key)
 
         state = await wait_or_fail(
@@ -203,11 +191,7 @@ async def test_text_sensor_raw_state(
 
         # Test 4: Prepend filter (StringRef-based)
         # "prefix " + "test" = "prefix test"
-        prepend_button = next(
-            (e for e in entities if "test_prepend_button" in e.object_id.lower()),
-            None,
-        )
-        assert prepend_button is not None, "Test Prepend Button not found"
+        prepend_button = require_entity(entities, "test_prepend_button")
         client.button_command(prepend_button.key)
 
         state = await wait_or_fail(
@@ -220,11 +204,7 @@ async def test_text_sensor_raw_state(
 
         # Test 5: Substitute filter (StringRef-based)
         # "foo says hello" with foo->bar, hello->world = "bar says world"
-        substitute_button = next(
-            (e for e in entities if "test_substitute_button" in e.object_id.lower()),
-            None,
-        )
-        assert substitute_button is not None, "Test Substitute Button not found"
+        substitute_button = require_entity(entities, "test_substitute_button")
         client.button_command(substitute_button.key)
 
         state = await wait_or_fail(
@@ -236,11 +216,7 @@ async def test_text_sensor_raw_state(
         )
 
         # Test 6: Map filter - "ON" -> "Active"
-        map_on_button = next(
-            (e for e in entities if "test_map_on_button" in e.object_id.lower()),
-            None,
-        )
-        assert map_on_button is not None, "Test Map ON Button not found"
+        map_on_button = require_entity(entities, "test_map_on_button")
         client.button_command(map_on_button.key)
 
         state = await wait_or_fail(
@@ -250,11 +226,7 @@ async def test_text_sensor_raw_state(
         assert state == "Active", f"Map ON failed: expected 'Active', got '{state}'"
 
         # Test 7: Map filter - "OFF" -> "Inactive"
-        map_off_button = next(
-            (e for e in entities if "test_map_off_button" in e.object_id.lower()),
-            None,
-        )
-        assert map_off_button is not None, "Test Map OFF Button not found"
+        map_off_button = require_entity(entities, "test_map_off_button")
         client.button_command(map_off_button.key)
 
         state = await wait_or_fail(
@@ -267,11 +239,7 @@ async def test_text_sensor_raw_state(
 
         # Test 8: Map filter - passthrough for unknown values
         # "UNKNOWN" -> "UNKNOWN" (no match, passes through unchanged)
-        map_unknown_button = next(
-            (e for e in entities if "test_map_unknown_button" in e.object_id.lower()),
-            None,
-        )
-        assert map_unknown_button is not None, "Test Map Unknown Button not found"
+        map_unknown_button = require_entity(entities, "test_map_unknown_button")
         client.button_command(map_unknown_button.key)
 
         state = await wait_or_fail(
@@ -284,11 +252,7 @@ async def test_text_sensor_raw_state(
 
         # Test 9: Chained filters (prepend "[" + append "]")
         # "[" + "value" + "]" = "[value]"
-        chained_button = next(
-            (e for e in entities if "test_chained_button" in e.object_id.lower()),
-            None,
-        )
-        assert chained_button is not None, "Test Chained Button not found"
+        chained_button = require_entity(entities, "test_chained_button")
         client.button_command(chained_button.key)
 
         state = await wait_or_fail(
@@ -299,11 +263,7 @@ async def test_text_sensor_raw_state(
 
         # Test 10: to_lower filter
         # "HELLO WORLD" -> "hello world"
-        to_lower_button = next(
-            (e for e in entities if "test_to_lower_button" in e.object_id.lower()),
-            None,
-        )
-        assert to_lower_button is not None, "Test To Lower Button not found"
+        to_lower_button = require_entity(entities, "test_to_lower_button")
         client.button_command(to_lower_button.key)
 
         state = await wait_or_fail(
@@ -316,11 +276,7 @@ async def test_text_sensor_raw_state(
 
         # Test 11: Lambda filter
         # "test" -> "[test]"
-        lambda_button = next(
-            (e for e in entities if "test_lambda_button" in e.object_id.lower()),
-            None,
-        )
-        assert lambda_button is not None, "Test Lambda Button not found"
+        lambda_button = require_entity(entities, "test_lambda_button")
         client.button_command(lambda_button.key)
 
         state = await wait_or_fail(
@@ -331,11 +287,7 @@ async def test_text_sensor_raw_state(
 
         # Test 12: Lambda filter - value passes through
         # "value" -> "value passed"
-        lambda_pass_button = next(
-            (e for e in entities if "test_lambda_pass_button" in e.object_id.lower()),
-            None,
-        )
-        assert lambda_pass_button is not None, "Test Lambda Pass Button not found"
+        lambda_pass_button = require_entity(entities, "test_lambda_pass_button")
         client.button_command(lambda_pass_button.key)
 
         state = await wait_or_fail(
@@ -348,11 +300,7 @@ async def test_text_sensor_raw_state(
 
         # Test 13: Lambda filter - skip publishing (return {})
         # "skip" -> no publish, state remains "value passed" from previous test
-        lambda_skip_button = next(
-            (e for e in entities if "test_lambda_skip_button" in e.object_id.lower()),
-            None,
-        )
-        assert lambda_skip_button is not None, "Test Lambda Skip Button not found"
+        lambda_skip_button = require_entity(entities, "test_lambda_skip_button")
         client.button_command(lambda_skip_button.key)
 
         state = await wait_or_fail(
@@ -368,16 +316,8 @@ async def test_text_sensor_raw_state(
         # Test 14: Lambda filter - verify raw_state is preserved (not mutated)
         # This is critical to verify the in-place mutation optimization is safe
         # "original" -> state="original MODIFIED", raw_state="original"
-        lambda_raw_state_button = next(
-            (
-                e
-                for e in entities
-                if "test_lambda_raw_state_button" in e.object_id.lower()
-            ),
-            None,
-        )
-        assert lambda_raw_state_button is not None, (
-            "Test Lambda Raw State Button not found"
+        lambda_raw_state_button = require_entity(
+            entities, "test_lambda_raw_state_button"
         )
         client.button_command(lambda_raw_state_button.key)
 

tests/integration/test_water_heater_template.py

@@ -14,7 +14,7 @@ from aioesphomeapi import (
 )
 import pytest
 
-from .state_utils import InitialStateHelper
+from .state_utils import InitialStateHelper, require_only_entity
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -51,12 +51,7 @@ async def test_water_heater_template(
         # Get entities and set up state synchronization
         entities, services = await client.list_entities_services()
         initial_state_helper = InitialStateHelper(entities)
-        water_heater_infos = [e for e in entities if isinstance(e, WaterHeaterInfo)]
-        assert len(water_heater_infos) == 1, (
-            f"Expected exactly 1 water heater entity, got {len(water_heater_infos)}. Entity types: {[type(e).__name__ for e in entities]}"
-        )
-
-        test_water_heater = water_heater_infos[0]
+        test_water_heater = require_only_entity(entities, WaterHeaterInfo)
 
         # Verify water heater entity info
         assert test_water_heater.object_id == "test_boiler"
@@ -177,8 +172,7 @@ async def test_water_heater_template_unknown_temperature(
                 state_count += 1
 
         entities, _ = await client.list_entities_services()
-        water_heater_infos = [e for e in entities if isinstance(e, WaterHeaterInfo)]
-        assert len(water_heater_infos) == 1
+        require_only_entity(entities, WaterHeaterInfo)
 
         client.subscribe_states(on_state)