From f47dca7e4eead77fe08cc029fc8384f1e1f264a6 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Tue, 20 Oct 2026 15:42:31 -1000
Subject: [PATCH] [tests] Use the exact entity lookups in the integration tests
 (#21213)

---
 tests/integration/state_utils.py              |  4 +-
 .../test_api_message_size_batching.py         | 73 +++------------
 ...test_batch_delay_zero_rapid_transitions.py |  4 +-
 tests/integration/test_build_info.py          | 18 +---
 .../integration/test_climate_custom_modes.py  |  7 +-
 .../test_large_message_batching.py            | 12 +--
 .../integration/test_legacy_climate_compat.py |  9 +-
 tests/integration/test_legacy_fan_compat.py   |  6 +-
 tests/integration/test_light_automations.py   |  3 +-
 .../test_light_binary_effect_off_phase.py     |  6 +-
 .../integration/test_light_control_action.py  |  3 +-
 .../test_light_effect_zero_brightness.py      |  4 +-
 tests/integration/test_light_repeat_effect.py |  4 +-
 tests/integration/test_light_resume_effect.py |  8 +-
 tests/integration/test_lock_automations.py    |  3 +-
 .../test_scheduler_defer_cancel.py            | 20 ++--
 .../test_scheduler_defer_fifo_simple.py       | 20 ++--
 .../test_sensor_filters_ring_buffer.py        |  8 +-
 .../test_sensor_filters_sliding_window.py     | 20 +---
 tests/integration/test_sensor_raw_state.py    |  6 +-
 .../integration/test_sensor_timeout_filter.py |  9 +-
 .../test_template_climate_basic.py            | 11 ++-
 .../test_template_climate_custom_modes.py     | 11 ++-
 .../test_template_climate_nonoptimistic.py    | 11 ++-
 ...st_template_climate_on_control_ordering.py |  6 +-
 ...est_template_climate_publish_all_fields.py | 10 +-
 .../test_template_climate_sensor_push.py      | 11 ++-
 .../test_template_climate_set_actions.py      | 11 ++-
 ..._template_climate_two_point_temperature.py | 12 ++-
 .../integration/test_text_sensor_raw_state.py | 92 ++++---------------
 .../integration/test_water_heater_template.py | 12 +--
 31 files changed, 142 insertions(+), 292 deletions(-)

diff --git a/tests/integration/state_utils.py b/tests/integration/state_utils.py
index 909164973e2..3c0e77cd161 100644
--- a/tests/integration/state_utils.py
+++ b/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
 
 
diff --git a/tests/integration/test_api_message_size_batching.py b/tests/integration/test_api_message_size_batching.py
index 626ac93ba62..51ce639969e 100644
--- a/tests/integration/test_api_message_size_batching.py
+++ b/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}"
         )
diff --git a/tests/integration/test_batch_delay_zero_rapid_transitions.py b/tests/integration/test_batch_delay_zero_rapid_transitions.py
index f17319dddfc..b68a3b49378 100644
--- a/tests/integration/test_batch_delay_zero_rapid_transitions.py
+++ b/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)
diff --git a/tests/integration/test_build_info.py b/tests/integration/test_build_info.py
index e25f7ab340b..356a2e52ada 100644
--- a/tests/integration/test_build_info.py
+++ b/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()
diff --git a/tests/integration/test_climate_custom_modes.py b/tests/integration/test_climate_custom_modes.py
index ff9d7476195..45fda95ef89 100644
--- a/tests/integration/test_climate_custom_modes.py
+++ b/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))
diff --git a/tests/integration/test_large_message_batching.py b/tests/integration/test_large_message_batching.py
index c80c6c25329..f91de21378c 100644
--- a/tests/integration/test_large_message_batching.py
+++ b/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
diff --git a/tests/integration/test_legacy_climate_compat.py b/tests/integration/test_legacy_climate_compat.py
index db37be3974c..fa09c6b0894 100644
--- a/tests/integration/test_legacy_climate_compat.py
+++ b/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) == {
diff --git a/tests/integration/test_legacy_fan_compat.py b/tests/integration/test_legacy_fan_compat.py
index 0825a7459b8..276dd2bfbaa 100644
--- a/tests/integration/test_legacy_fan_compat.py
+++ b/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) == {
diff --git a/tests/integration/test_light_automations.py b/tests/integration/test_light_automations.py
index 86376df9b88..7a083a48734 100644
--- a/tests/integration/test_light_automations.py
+++ b/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
diff --git a/tests/integration/test_light_binary_effect_off_phase.py b/tests/integration/test_light_binary_effect_off_phase.py
index 571fca19f65..26e270dfaf1 100644
--- a/tests/integration/test_light_binary_effect_off_phase.py
+++ b/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]] = {}
 
diff --git a/tests/integration/test_light_control_action.py b/tests/integration/test_light_control_action.py
index bb690645bd4..89a98508949 100644
--- a/tests/integration/test_light_control_action.py
+++ b/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:
diff --git a/tests/integration/test_light_effect_zero_brightness.py b/tests/integration/test_light_effect_zero_brightness.py
index 6c386d4229d..4581ef8a023 100644
--- a/tests/integration/test_light_effect_zero_brightness.py
+++ b/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]] = {}
 
diff --git a/tests/integration/test_light_repeat_effect.py b/tests/integration/test_light_repeat_effect.py
index e8e4cf8e12d..6869b213fe6 100644
--- a/tests/integration/test_light_repeat_effect.py
+++ b/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]] = {}
 
diff --git a/tests/integration/test_light_resume_effect.py b/tests/integration/test_light_resume_effect.py
index dce5feb9ab5..ebf1882becd 100644
--- a/tests/integration/test_light_resume_effect.py
+++ b/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]] = {}
 
diff --git a/tests/integration/test_lock_automations.py b/tests/integration/test_lock_automations.py
index a1098164d7c..d7994168f09 100644
--- a/tests/integration/test_lock_automations.py
+++ b/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)
diff --git a/tests/integration/test_scheduler_defer_cancel.py b/tests/integration/test_scheduler_defer_cancel.py
index 83f706b7726..9a3d8c201cd 100644
--- a/tests/integration/test_scheduler_defer_cancel.py
+++ b/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
diff --git a/tests/integration/test_scheduler_defer_fifo_simple.py b/tests/integration/test_scheduler_defer_fifo_simple.py
index 39a4e0ccbed..ff068265a95 100644
--- a/tests/integration/test_scheduler_defer_fifo_simple.py
+++ b/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
diff --git a/tests/integration/test_sensor_filters_ring_buffer.py b/tests/integration/test_sensor_filters_ring_buffer.py
index 5f0fac6623e..ad8d59ac906 100644
--- a/tests/integration/test_sensor_filters_ring_buffer.py
+++ b/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)
diff --git a/tests/integration/test_sensor_filters_sliding_window.py b/tests/integration/test_sensor_filters_sliding_window.py
index d873898c120..022f99068d8 100644
--- a/tests/integration/test_sensor_filters_sliding_window.py
+++ b/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
diff --git a/tests/integration/test_sensor_raw_state.py b/tests/integration/test_sensor_raw_state.py
index f354690fcee..d0d2abe4f86 100644
--- a/tests/integration/test_sensor_raw_state.py
+++ b/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")
 
diff --git a/tests/integration/test_sensor_timeout_filter.py b/tests/integration/test_sensor_timeout_filter.py
index 4c5a95daac2..edbd62a223b 100644
--- a/tests/integration/test_sensor_timeout_filter.py
+++ b/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")
diff --git a/tests/integration/test_template_climate_basic.py b/tests/integration/test_template_climate_basic.py
index 431fd4e3e87..e68966d3b0a 100644
--- a/tests/integration/test_template_climate_basic.py
+++ b/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
diff --git a/tests/integration/test_template_climate_custom_modes.py b/tests/integration/test_template_climate_custom_modes.py
index 4817fe1ddfc..01da12feccb 100644
--- a/tests/integration/test_template_climate_custom_modes.py
+++ b/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",
diff --git a/tests/integration/test_template_climate_nonoptimistic.py b/tests/integration/test_template_climate_nonoptimistic.py
index e922ec31b9c..445775212f7 100644
--- a/tests/integration/test_template_climate_nonoptimistic.py
+++ b/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
diff --git a/tests/integration/test_template_climate_on_control_ordering.py b/tests/integration/test_template_climate_on_control_ordering.py
index 8d212b3ccba..3d5dcdd8579 100644
--- a/tests/integration/test_template_climate_on_control_ordering.py
+++ b/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)
diff --git a/tests/integration/test_template_climate_publish_all_fields.py b/tests/integration/test_template_climate_publish_all_fields.py
index 9c4262b3113..1c06d410f90 100644
--- a/tests/integration/test_template_climate_publish_all_fields.py
+++ b/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)
 
diff --git a/tests/integration/test_template_climate_sensor_push.py b/tests/integration/test_template_climate_sensor_push.py
index 1db4da81edb..2c253676fa8 100644
--- a/tests/integration/test_template_climate_sensor_push.py
+++ b/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(
diff --git a/tests/integration/test_template_climate_set_actions.py b/tests/integration/test_template_climate_set_actions.py
index 0b1eb808740..3e16fe2c727 100644
--- a/tests/integration/test_template_climate_set_actions.py
+++ b/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(
diff --git a/tests/integration/test_template_climate_two_point_temperature.py b/tests/integration/test_template_climate_two_point_temperature.py
index 9270b59ffc1..cc82f52d0a4 100644
--- a/tests/integration/test_template_climate_two_point_temperature.py
+++ b/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
diff --git a/tests/integration/test_text_sensor_raw_state.py b/tests/integration/test_text_sensor_raw_state.py
index 3a5010f6fbb..313ab2afe45 100644
--- a/tests/integration/test_text_sensor_raw_state.py
+++ b/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)
 
diff --git a/tests/integration/test_water_heater_template.py b/tests/integration/test_water_heater_template.py
index 3d7f8851605..86671a896da 100644
--- a/tests/integration/test_water_heater_template.py
+++ b/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)