[tests] Use require_entity and require_only_entity for the remaining entity lookups (#21243)

tests/integration/state_utils.py

@@ -622,7 +622,6 @@ class SensorTracker:
         if match_initial_states:
             for state in initial_state_helper.initial_states.values():
                 self.on_state(state, first_pending_only=True)
-        start_btn = find_entity(entities, "start_scenario", ButtonInfo)
-        assert start_btn is not None, "Start Scenario button not found"
+        start_btn = require_entity(entities, "start_scenario", ButtonInfo)
         client.button_command(start_btn.key)
         return entities

tests/integration/test_alarm_control_panel_state_transitions.py

@@ -15,7 +15,7 @@ from aioesphomeapi import (
 )
 import pytest
 
-from .state_utils import InitialStateHelper
+from .state_utils import InitialStateHelper, require_entity, require_only_entity
 from .types import APIClientConnectedFactory, RunCompiledFunction
 from .wait_utils import wait_or_fail
 
@@ -79,23 +79,13 @@ async def test_alarm_control_panel_state_transitions(
     ):
         entities, _ = await client.list_entities_services()
 
-        # Find entities
-        alarm_info: AlarmControlPanelInfo | None = None
-        door_switch_info: SwitchInfo | None = None
-        chime_switch_info: SwitchInfo | None = None
-
-        for entity in entities:
-            if isinstance(entity, AlarmControlPanelInfo):
-                alarm_info = entity
-            elif isinstance(entity, SwitchInfo):
-                if entity.name == "Door Sensor Switch":
-                    door_switch_info = entity
-                elif entity.name == "Chime Sensor Switch":
-                    chime_switch_info = entity
-
-        assert alarm_info is not None, "Alarm control panel not found"
-        assert door_switch_info is not None, "Door sensor switch not found"
-        assert chime_switch_info is not None, "Chime sensor switch not found"
+        alarm_info = require_only_entity(entities, AlarmControlPanelInfo)
+        door_switch_info = require_entity(
+            entities, "door_sensor_switch", SwitchInfo, exact=True
+        )
+        chime_switch_info = require_entity(
+            entities, "chime_sensor_switch", SwitchInfo, exact=True
+        )
 
         # Track state changes
         states_received: list[AlarmControlPanelState] = []

tests/integration/test_api_homeassistant_action_no_subscriber.py

@@ -13,7 +13,7 @@ import asyncio
 from aioesphomeapi import ButtonInfo, HomeassistantServiceCall
 import pytest
 
-from .state_utils import assert_device_name
+from .state_utils import assert_device_name, require_only_entity
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -71,7 +71,7 @@ async def test_api_homeassistant_action_no_subscriber(
             client.subscribe_service_calls(on_service_call)
 
             entities, _ = await client.list_entities_services()
-            button = next(e for e in entities if isinstance(e, ButtonInfo))
+            button = require_only_entity(entities, ButtonInfo)
             client.button_command(button.key)
 
             button_call = await asyncio.wait_for(button_action_future, timeout=5.0)

tests/integration/test_host_logger_thread_safety.py

@@ -11,7 +11,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
 
@@ -62,11 +62,8 @@ async def test_host_logger_thread_safety(
         # Verify connection works
         await assert_device_name(client, "host-logger-thread-test")
 
-        # Get the button entity - find by name
         entities, _ = await client.list_entities_services()
-        button_entities = [e for e in entities if e.name == "Start Thread Race Test"]
-        assert button_entities, "Could not find Start Thread Race Test button"
-        button_key = button_entities[0].key
+        button_key = require_entity(entities, "start_thread_race_test", exact=True).key
 
         # Press the button to start the thread race test
         client.button_command(button_key)

tests/integration/test_host_mode_climate_basic_state.py

@@ -12,7 +12,7 @@ from aioesphomeapi import (
 )
 import pytest
 
-from .state_utils import wait_for_state
+from .state_utils import require_only_entity, wait_for_state
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -25,9 +25,7 @@ async def test_host_mode_climate_basic_state(
     """Test basic climate state reporting."""
     async with run_compiled(yaml_config), api_client_connected() as client:
         entities, _ = await client.list_entities_services()
-        climate_infos = [e for e in entities if isinstance(e, ClimateInfo)]
-        assert len(climate_infos) >= 1, "Expected at least 1 climate entity"
-        test_climate = climate_infos[0]
+        test_climate = require_only_entity(entities, ClimateInfo)
 
         # The thermostat publishes multiple states during setup as the
         # temperature/humidity sensors come online. Wait for the state to

tests/integration/test_host_mode_climate_control.py

@@ -8,7 +8,7 @@ import aioesphomeapi
 from aioesphomeapi import ClimateInfo, ClimateMode, 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
 
@@ -39,8 +39,8 @@ async def test_host_mode_climate_control(
         # 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 at least 1 climate entity"
+        test_climate = require_only_entity(entities, ClimateInfo)
+        assert test_climate.name == "Dual-mode Thermostat"
 
         # Subscribe with the wrapper that filters initial states
         client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
@@ -51,13 +51,6 @@ async def test_host_mode_climate_control(
         except TimeoutError:
             pytest.fail("Timeout waiting for initial states")
 
-        test_climate = next(
-            (c for c in climate_infos if c.name == "Dual-mode Thermostat"), None
-        )
-        assert test_climate is not None, (
-            "Dual-mode Thermostat thermostat climate not found"
-        )
-
         # Adjust setpoints
         client.climate_command(
             test_climate.key,

tests/integration/test_host_mode_many_entities.py

@@ -20,6 +20,7 @@ from aioesphomeapi import (
 )
 import pytest
 
+from .state_utils import require_only_entity
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -120,10 +121,7 @@ async def test_host_mode_many_entities(
 
         # Get entity info to verify climate entity details
         entities = await client.list_entities_services()
-        climate_infos = [e for e in entities[0] if isinstance(e, ClimateInfo)]
-        assert len(climate_infos) >= 1, "Expected at least 1 climate entity"
-
-        climate_info = climate_infos[0]
+        climate_info = require_only_entity(entities[0], ClimateInfo)
 
         # Verify feature flags set as expected
         assert climate_info.feature_flags == (
@@ -177,27 +175,18 @@ async def test_host_mode_many_entities(
         assert "AWAY" in preset_names, f"Expected 'AWAY' preset, got {preset_names}"
         assert "SLEEP" in preset_names, f"Expected 'SLEEP' preset, got {preset_names}"
 
-        # Verify datetime entities exist
-        date_infos = [e for e in entities[0] if isinstance(e, DateInfo)]
-        time_infos = [e for e in entities[0] if isinstance(e, TimeInfo)]
-        datetime_infos = [e for e in entities[0] if isinstance(e, DateTimeInfo)]
-
-        assert len(date_infos) >= 1, "Expected at least 1 date entity"
-        assert len(time_infos) >= 1, "Expected at least 1 time entity"
-        assert len(datetime_infos) >= 1, "Expected at least 1 datetime entity"
-
-        # Verify the entity names
-        date_info = date_infos[0]
+        # Verify the datetime entities and their names
+        date_info = require_only_entity(entities[0], DateInfo)
         assert date_info.name == "Test Date", (
             f"Expected date entity name 'Test Date', got {date_info.name}"
         )
 
-        time_info = time_infos[0]
+        time_info = require_only_entity(entities[0], TimeInfo)
         assert time_info.name == "Test Time", (
             f"Expected time entity name 'Test Time', got {time_info.name}"
         )
 
-        datetime_info = datetime_infos[0]
+        datetime_info = require_only_entity(entities[0], DateTimeInfo)
         assert datetime_info.name == "Test DateTime", (
             f"Expected datetime entity name 'Test DateTime', got {datetime_info.name}"
         )

tests/integration/test_host_preferences.py

@@ -4,24 +4,15 @@ from __future__ import annotations
 
 import asyncio
 import re
-from typing import Any
 
-from aioesphomeapi import ButtonInfo, EntityInfo, NumberInfo, SwitchInfo
+from aioesphomeapi import ButtonInfo, NumberInfo, SwitchInfo
 import pytest
 
+from .state_utils import require_entity
 from .types import APIClientConnectedFactory, RunCompiledFunction
 from .wait_utils import wait_or_fail
 
 
-def find_entity_by_name(
-    entities: list[EntityInfo], entity_type: type, name: str
-) -> Any:
-    """Helper to find an entity by type and name."""
-    return next(
-        (e for e in entities if isinstance(e, entity_type) and e.name == name), None
-    )
-
-
 @pytest.mark.asyncio
 async def test_host_preferences_save_load(
     yaml_config: str,
@@ -84,18 +75,17 @@ async def test_host_preferences_save_load(
         # Get entity list
         entities, _ = await client.list_entities_services()
 
-        # Find our test entities using helper
-        test_switch = find_entity_by_name(entities, SwitchInfo, "Test Switch")
-        test_number = find_entity_by_name(entities, NumberInfo, "Test Number")
-        save_button = find_entity_by_name(entities, ButtonInfo, "Save Preferences")
-        load_button = find_entity_by_name(entities, ButtonInfo, "Load Preferences")
-        verify_button = find_entity_by_name(entities, ButtonInfo, "Verify Preferences")
-
-        assert test_switch is not None, "Test Switch not found"
-        assert test_number is not None, "Test Number not found"
-        assert save_button is not None, "Save Preferences button not found"
-        assert load_button is not None, "Load Preferences button not found"
-        assert verify_button is not None, "Verify Preferences button not found"
+        test_switch = require_entity(entities, "test_switch", SwitchInfo, exact=True)
+        test_number = require_entity(entities, "test_number", NumberInfo, exact=True)
+        save_button = require_entity(
+            entities, "save_preferences", ButtonInfo, exact=True
+        )
+        load_button = require_entity(
+            entities, "load_preferences", ButtonInfo, exact=True
+        )
+        verify_button = require_entity(
+            entities, "verify_preferences", ButtonInfo, exact=True
+        )
 
         # Set initial values
         client.switch_command(test_switch.key, True)

tests/integration/test_host_preferences_suspend_resume.py

@@ -5,26 +5,17 @@ from __future__ import annotations
 import asyncio
 from pathlib import Path
 import re
-from typing import Any
 
-from aioesphomeapi import ButtonInfo, EntityInfo
+from aioesphomeapi import ButtonInfo
 import pytest
 
+from .state_utils import require_entity
 from .types import APIClientConnectedFactory, RunCompiledFunction
 from .wait_utils import wait_or_fail
 
 DEVICE_NAME = "test_suspend_resume_device"
 
 
-def find_entity_by_name(
-    entities: list[EntityInfo], entity_type: type, name: str
-) -> Any:
-    """Helper to find an entity by type and name."""
-    return next(
-        (e for e in entities if isinstance(e, entity_type) and e.name == name), None
-    )
-
-
 async def _poll_until_exists(path: Path) -> None:
     """Poll for a file to appear, rather than guessing a sleep duration."""
     while not path.exists():
@@ -64,12 +55,15 @@ async def test_host_preferences_suspend_resume(
     ):
         entities, _ = await client.list_entities_services()
 
-        save_button = find_entity_by_name(entities, ButtonInfo, "Save Preference")
-        suspend_button = find_entity_by_name(entities, ButtonInfo, "Suspend Syncer")
-        resume_button = find_entity_by_name(entities, ButtonInfo, "Resume Syncer")
-        assert save_button is not None, "Save Preference button not found"
-        assert suspend_button is not None, "Suspend Syncer button not found"
-        assert resume_button is not None, "Resume Syncer button not found"
+        save_button = require_entity(
+            entities, "save_preference", ButtonInfo, exact=True
+        )
+        suspend_button = require_entity(
+            entities, "suspend_syncer", ButtonInfo, exact=True
+        )
+        resume_button = require_entity(
+            entities, "resume_syncer", ButtonInfo, exact=True
+        )
 
         # --- Positive control: a running syncer flushes to disk. Without this,
         # the suspend assertion below could pass for the wrong reason (e.g. wrong prefs path). ---

tests/integration/test_ir_rf_transmit_complete.py

@@ -15,7 +15,7 @@ from aioesphomeapi import InfraredInfo, RadioFrequencyInfo
 from aioesphomeapi.api_pb2 import InfraredRFTransmitCompleteResponse
 import pytest
 
-from .state_utils import find_entity
+from .state_utils import require_entity
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 FRAME_COUNT = 5
@@ -37,9 +37,9 @@ async def test_ir_rf_transmit_complete_boot(
     """The host build with the mock transmitters boots and lists all entities on any client."""
     async with run_compiled(yaml_config), api_client_connected() as client:
         entities, _ = await client.list_entities_services()
-    assert find_entity(entities, "rf_transmitter_a", RadioFrequencyInfo) is not None
-    assert find_entity(entities, "rf_transmitter_b", RadioFrequencyInfo) is not None
-    assert find_entity(entities, "ir_transmitter", InfraredInfo) is not None
+    require_entity(entities, "rf_transmitter_a", RadioFrequencyInfo)
+    require_entity(entities, "rf_transmitter_b", RadioFrequencyInfo)
+    require_entity(entities, "ir_transmitter", InfraredInfo)
 
 
 @pytest.mark.shared_yaml("ir_rf_transmit_complete")
@@ -92,11 +92,9 @@ async def test_ir_rf_transmit_complete(
         api_client_connected() as client,
     ):
         entities, _ = await client.list_entities_services()
-        rf = find_entity(entities, "rf_transmitter_a", RadioFrequencyInfo)
-        rf_b = find_entity(entities, "rf_transmitter_b", RadioFrequencyInfo)
-        ir = find_entity(entities, "ir_transmitter", InfraredInfo)
-        assert rf is not None and rf_b is not None, "RF transmitter entities not found"
-        assert ir is not None, "IR transmitter entity not found"
+        rf = require_entity(entities, "rf_transmitter_a", RadioFrequencyInfo)
+        rf_b = require_entity(entities, "rf_transmitter_b", RadioFrequencyInfo)
+        ir = require_entity(entities, "ir_transmitter", InfraredInfo)
 
         client._connection.add_message_callback(
             on_complete, (InfraredRFTransmitCompleteResponse,)

tests/integration/test_light_constant_brightness.py

@@ -19,7 +19,7 @@ from typing import Any
 from aioesphomeapi import EntityState, LightInfo, LightState
 import pytest
 
-from .state_utils import InitialStateHelper
+from .state_utils import InitialStateHelper, require_entity
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -61,9 +61,8 @@ async def test_light_constant_brightness(
         api_client_connected() as client,
     ):
         entities, _ = await client.list_entities_services()
-        lights = [e for e in entities if isinstance(e, LightInfo)]
-        cb_light = next(e for e in lights if e.object_id.endswith("cb_light"))
-        ncb_light = next(e for e in lights if e.object_id.endswith("ncb_light"))
+        cb_light = require_entity(entities, "test_cb_light", LightInfo, exact=True)
+        ncb_light = require_entity(entities, "test_ncb_light", LightInfo, exact=True)
 
         # Use InitialStateHelper to wait for initial state broadcast
         initial_state_helper = InitialStateHelper(entities)

tests/integration/test_logger_buffered_recursion_guard.py

@@ -29,7 +29,7 @@ import re
 from aioesphomeapi import LogLevel
 import pytest
 
-from .state_utils import assert_device_name
+from .state_utils import assert_device_name, require_entity
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 _ANSI = re.compile(r"\x1b\[[0-9;]*m")
@@ -73,9 +73,8 @@ async def test_logger_buffered_recursion_guard(
         client.subscribe_logs(on_log, log_level=LogLevel.LOG_LEVEL_VERY_VERBOSE)
 
         entities, _ = await client.list_entities_services()
-        buttons = [e for e in entities if e.name == "Start Race Test"]
-        assert buttons, "Could not find Start Race Test button"
-        client.button_command(buttons[0].key)
+        button = require_entity(entities, "start_race_test", exact=True)
+        client.button_command(button.key)
 
         # Wait until every buffered thread message has been delivered over the API.
         try:

tests/integration/test_multi_device_preferences.py

@@ -8,6 +8,7 @@ import re
 from aioesphomeapi import ButtonInfo, NumberInfo, SelectInfo, SwitchInfo
 import pytest
 
+from .state_utils import require_entity
 from .types import APIClientConnectedFactory, RunCompiledFunction
 from .wait_utils import wait_or_fail
 
@@ -97,9 +98,9 @@ async def test_multi_device_preferences(
         assert len(selects) == 3, f"Expected 3 'Mode' selects, got {len(selects)}"
 
         # Find the test button entity to trigger preference logging
-        buttons = [e for e in entities if isinstance(e, ButtonInfo)]
-        test_button = next((b for b in buttons if b.name == "Test Preferences"), None)
-        assert test_button is not None, "Test Preferences button not found"
+        test_button = require_entity(
+            entities, "test_preferences", ButtonInfo, exact=True
+        )
 
         # Press the button to trigger logging
         client.button_command(test_button.key)

tests/integration/test_select_stringref_trigger.py

@@ -7,7 +7,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
 
 
@@ -95,14 +95,11 @@ async def test_select_stringref_trigger(
         # List entities to find our select
         entities, _ = await client.list_entities_services()
 
-        select_entity = next((e for e in entities if e.name == "Test Select"), None)
-        assert select_entity is not None, "Test Select entity not found"
+        select_entity = require_entity(entities, "test_select", exact=True)
 
-        baud_entity = next((e for e in entities if e.name == "Baud Rate"), None)
-        assert baud_entity is not None, "Baud Rate entity not found"
+        baud_entity = require_entity(entities, "baud_rate", exact=True)
 
-        action_entity = next((e for e in entities if e.name == "Action Select"), None)
-        assert action_entity is not None, "Action Select entity not found"
+        action_entity = require_entity(entities, "action_select", exact=True)
 
         # Change select to Option B - this should trigger on_value with StringRef
         client.select_command(select_entity.key, "Option B")

tests/integration/test_sensor_timeout_throttle_filter.py

@@ -8,7 +8,7 @@ import math
 from aioesphomeapi import ButtonInfo, 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
 
@@ -43,10 +43,8 @@ async def test_sensor_timeout_throttle_filter(
         client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
         await initial_state_helper.wait_for_initial_states()
 
-        button = next(
-            e.key for e in entities if isinstance(e, ButtonInfo) and e.name == "Drive"
-        )
-        client.button_command(button)
+        button = require_entity(entities, "drive", ButtonInfo, exact=True)
+        client.button_command(button.key)
         await wait_or_fail(
             asyncio.gather(*done.values()), 5.0, "Timed out", values=values
         )

tests/integration/test_template_alarm_control_panel_many_sensors.py

@@ -6,7 +6,7 @@ import aioesphomeapi
 from aioesphomeapi.model import APIIntEnum
 import pytest
 
-from .state_utils import InitialStateHelper
+from .state_utils import InitialStateHelper, require_only_entity
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -32,17 +32,12 @@ async def test_template_alarm_control_panel_many_sensors(
         # Get entity info first
         entities, _ = await client.list_entities_services()
 
-        # Find the alarm control panel and binary sensors
-        alarm_info: aioesphomeapi.AlarmControlPanelInfo | None = None
-        binary_sensors: list[aioesphomeapi.BinarySensorInfo] = []
-
-        for entity in entities:
-            if isinstance(entity, aioesphomeapi.AlarmControlPanelInfo):
-                alarm_info = entity
-            elif isinstance(entity, aioesphomeapi.BinarySensorInfo):
-                binary_sensors.append(entity)
-
-        assert alarm_info is not None, "Alarm control panel entity info not found"
+        alarm_info = require_only_entity(entities, aioesphomeapi.AlarmControlPanelInfo)
+        binary_sensors = [
+            entity
+            for entity in entities
+            if isinstance(entity, aioesphomeapi.BinarySensorInfo)
+        ]
         assert alarm_info.name == "Test Alarm"
         assert alarm_info.requires_code is True
         assert alarm_info.requires_code_to_arm is True

tests/integration/test_uart_mock_ld2410.py

@@ -23,7 +23,7 @@ import asyncio
 from aioesphomeapi import ButtonInfo
 import pytest
 
-from .state_utils import InitialStateHelper, SensorStateCollector, find_entity
+from .state_utils import InitialStateHelper, SensorStateCollector, require_entity
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -95,8 +95,7 @@ async def test_uart_mock_ld2410(
             pytest.fail("Timeout waiting for initial states")
 
         # Start the UART mock scenario now that we're subscribed
-        start_btn = find_entity(entities, "start_scenario", ButtonInfo)
-        assert start_btn is not None, "Start Scenario button not found"
+        start_btn = require_entity(entities, "start_scenario", ButtonInfo)
         client.button_command(start_btn.key)
 
         # Wait for Phase 1 - all sensors and binary sensors have at least one value
@@ -221,8 +220,7 @@ async def test_uart_mock_ld2410_engineering(
             pytest.fail("Timeout waiting for initial states")
 
         # Start the UART mock scenario now that we're subscribed
-        start_btn = find_entity(entities, "start_scenario", ButtonInfo)
-        assert start_btn is not None, "Start Scenario button not found"
+        start_btn = require_entity(entities, "start_scenario", ButtonInfo)
         client.button_command(start_btn.key)
 
         # Wait for Phase 1 - all sensors and binary sensors have at least one value

tests/integration/test_uart_mock_ld2410_query_ack.py

@@ -20,7 +20,7 @@ from aioesphomeapi import (
 )
 import pytest
 
-from .state_utils import InitialStateHelper, build_key_to_entity_mapping, find_entity
+from .state_utils import InitialStateHelper, build_key_to_entity_mapping, require_entity
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 # Number object_id: (value from the seed ACK, value from the ACK under test)
@@ -124,8 +124,7 @@ async def test_uart_mock_ld2410_query_ack(
             pytest.fail("Timeout waiting for initial states")
 
         # Start the UART mock scenario now that we're subscribed
-        start_btn = find_entity(entities, "start_scenario", ButtonInfo)
-        assert start_btn is not None, "Start Scenario button not found"
+        start_btn = require_entity(entities, "start_scenario", ButtonInfo)
         client.button_command(start_btn.key)
 
         try:

tests/integration/test_uart_mock_ld2412.py

@@ -46,7 +46,6 @@ from .state_utils import (
     InitialStateHelper,
     SensorStateCollector,
     StateWaiter,
-    find_entity,
     require_entity,
 )
 from .types import APIClientConnectedFactory, RunCompiledFunction
@@ -193,8 +192,7 @@ async def test_uart_mock_ld2412(
         await _wait_initial(client, entities, collector.on_state)
 
         # Start the UART mock scenario now that we're subscribed
-        start_btn = find_entity(entities, "start_scenario", ButtonInfo)
-        assert start_btn is not None, "Start Scenario button not found"
+        start_btn = require_entity(entities, "start_scenario", ButtonInfo)
         client.button_command(start_btn.key)
 
         # Wait for Phase 1 - all sensors and binary sensors have at least one value
@@ -301,8 +299,7 @@ async def test_uart_mock_ld2412_engineering(
         await _wait_initial(client, entities, collector.on_state)
 
         # Start the UART mock scenario now that we're subscribed
-        start_btn = find_entity(entities, "start_scenario", ButtonInfo)
-        assert start_btn is not None, "Start Scenario button not found"
+        start_btn = require_entity(entities, "start_scenario", ButtonInfo)
         client.button_command(start_btn.key)
 
         # Wait for Phase 1 - all sensors and binary sensors have at least one value
@@ -407,8 +404,7 @@ async def test_uart_mock_ld2412_engineering_truncated(
 
         await _wait_initial(client, entities, collector.on_state)
 
-        start_btn = find_entity(entities, "start_scenario", ButtonInfo)
-        assert start_btn is not None, "Start Scenario button not found"
+        start_btn = require_entity(entities, "start_scenario", ButtonInfo)
         client.button_command(start_btn.key)
 
         # Wait for Phase 1 — valid engineering frame establishes baseline

tests/integration/test_uart_mock_ld2420.py

@@ -24,7 +24,7 @@ import asyncio
 from aioesphomeapi import ButtonInfo
 import pytest
 
-from .state_utils import InitialStateHelper, SensorStateCollector, find_entity
+from .state_utils import InitialStateHelper, SensorStateCollector, require_entity
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -86,8 +86,7 @@ async def test_uart_mock_ld2420(
             pytest.fail("Timeout waiting for initial states")
 
         # Start the UART mock scenario now that we're subscribed
-        start_btn = find_entity(entities, "start_scenario", ButtonInfo)
-        assert start_btn is not None, "Start Scenario button not found"
+        start_btn = require_entity(entities, "start_scenario", ButtonInfo)
         client.button_command(start_btn.key)
 
         # Wait for Phase 1 - all sensors and binary sensors have at least one value
@@ -198,8 +197,7 @@ async def test_uart_mock_ld2420_simple(
             pytest.fail("Timeout waiting for initial states")
 
         # Start the UART mock scenario now that we're subscribed
-        start_btn = find_entity(entities, "start_scenario", ButtonInfo)
-        assert start_btn is not None, "Start Scenario button not found"
+        start_btn = require_entity(entities, "start_scenario", ButtonInfo)
         client.button_command(start_btn.key)
 
         # Wait for Phase 1 - all sensors and binary sensors have at least one value

tests/integration/test_uart_mock_ld2450.py

@@ -22,7 +22,7 @@ import asyncio
 from aioesphomeapi import ButtonInfo
 import pytest
 
-from .state_utils import InitialStateHelper, SensorStateCollector, find_entity
+from .state_utils import InitialStateHelper, SensorStateCollector, require_entity
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -108,8 +108,7 @@ async def test_uart_mock_ld2450(
             pytest.fail("Timeout waiting for initial states")
 
         # Start the UART mock scenario now that we're subscribed
-        start_btn = find_entity(entities, "start_scenario", ButtonInfo)
-        assert start_btn is not None, "Start Scenario button not found"
+        start_btn = require_entity(entities, "start_scenario", ButtonInfo)
         client.button_command(start_btn.key)
 
         # Wait for Phase 1 - all sensors and binary sensors have at least one value

tests/integration/test_uart_mock_ld2460.py

@@ -19,7 +19,7 @@ import asyncio
 from aioesphomeapi import ButtonInfo
 import pytest
 
-from .state_utils import InitialStateHelper, SensorStateCollector, find_entity
+from .state_utils import InitialStateHelper, SensorStateCollector, require_entity
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -89,8 +89,7 @@ async def test_uart_mock_ld2460(
         except TimeoutError:
             pytest.fail("Timeout waiting for initial states")
 
-        start_btn = find_entity(entities, "start_scenario", ButtonInfo)
-        assert start_btn is not None, "Start Scenario button not found"
+        start_btn = require_entity(entities, "start_scenario", ButtonInfo)
         client.button_command(start_btn.key)
 
         # Phase 1: the datasheet example, one target at (1.5, 2.3)

tests/integration/test_uart_mock_modbus.py

@@ -23,7 +23,7 @@ from collections.abc import Callable
 from aioesphomeapi import ButtonInfo, NumberInfo, SwitchInfo, TextSensorState
 import pytest
 
-from .state_utils import SensorTracker, find_entity, require_entity, wait_for_state
+from .state_utils import SensorTracker, require_entity, wait_for_state
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -832,10 +832,8 @@ async def test_uart_mock_modbus_continuous(
         # Recovery path: a live continuous poll that starts failing goes offline, and the next update()
         # re-arms it once the device answers again. Silence the server so the poll's reads time out; with
         # max_cmd_retries=1 and send_wait_time=100ms the device trips offline quickly and streaming stops.
-        silence = find_entity(entities, "silence_server", SwitchInfo)
-        assert silence is not None, "Silence Server switch not found"
-        start = find_entity(entities, "start_scenario", ButtonInfo)
-        assert start is not None, "Start Scenario button not found"
+        silence = require_entity(entities, "silence_server", SwitchInfo)
+        start = require_entity(entities, "start_scenario", ButtonInfo)
 
         client.switch_command(silence.key, True)
         await asyncio.sleep(1.0)  # let the poll fail and the device trip offline
@@ -890,8 +888,7 @@ async def test_uart_mock_modbus_offline(
         # Register the recovery expectations before waking the device so no update is missed.
         online_future = tracker.expect("link_state", 1)
         value_future = tracker.expect("reg", 259)
-        serve_btn = find_entity(entities, "serve", ButtonInfo)
-        assert serve_btn is not None, "Serve button not found"
+        serve_btn = require_entity(entities, "serve", ButtonInfo)
         client.button_command(serve_btn.key)
 
         # The next offline-cadence retry gets an answer: back online, value published.
@@ -929,8 +926,7 @@ async def test_uart_mock_modbus_fairness(
         # Stop polling so the counters settle to a final, stable value (state
         # coalescing means intermediate values may be skipped, but the final
         # value is always delivered once changes stop).
-        stop_btn = find_entity(entities, "stop_scenario", ButtonInfo)
-        assert stop_btn is not None, "Stop Scenario button not found"
+        stop_btn = require_entity(entities, "stop_scenario", ButtonInfo)
         client.button_command(stop_btn.key)
         await asyncio.sleep(0.5)
 
@@ -1050,8 +1046,7 @@ async def test_uart_mock_modbus_register_offset(
         entities = await tracker.setup_and_start_scenario(client)
         await tracker.await_all(initial, timeout=4.0)
 
-        switch = find_entity(entities, "offset_switch", SwitchInfo)
-        assert switch is not None, "offset_switch not found"
+        switch = require_entity(entities, "offset_switch", SwitchInfo)
         client.switch_command(switch.key, True)
 
         # reg_11 (0x10 + offset 2/2) must receive the write; if the write went to 0x12 this times out.
@@ -1066,8 +1061,7 @@ async def test_uart_mock_modbus_register_offset(
         # into its address (0x10 + 6/2 = 0x13, residual byte 0) and joins the 0x10..0x13 range, so the
         # read lands in-bounds on 0xABCD (bit 0 set) -> ON. Pre-fix the whole byte offset folded into the
         # address (0x16); the server answers ILLEGAL_DATA_ADDRESS there and the switch never publishes.
-        read_switch = find_entity(entities, "read_offset_switch", SwitchInfo)
-        assert read_switch is not None, "read_offset_switch not found"
+        read_switch = require_entity(entities, "read_offset_switch", SwitchInfo)
         # The ON transition happened at the first poll and switch states are deduped, so this relies on
         # wait_for_state's fresh subscribe_states re-dumping every entity's current state.
         await wait_for_state(
@@ -1107,8 +1101,7 @@ async def test_uart_mock_modbus_lambda_write(
         entities = await tracker.setup_and_start_scenario(client)
         await tracker.await_change(initial, "reg_30", timeout=4.0)
 
-        switch = find_entity(entities, "cross_switch", SwitchInfo)
-        assert switch is not None, "cross_switch not found"
+        switch = require_entity(entities, "cross_switch", SwitchInfo)
         client.switch_command(switch.key, True)
 
         # The coil switch's lambda wrote register 0x30 via item->write_single_register(); reg_30 must
@@ -1144,8 +1137,7 @@ async def test_uart_mock_modbus_lambda_invert(
         entities = await tracker.setup_and_start_scenario(client)
         await tracker.await_change(initial, "reg_40", timeout=4.0)
 
-        switch = find_entity(entities, "invert_switch", SwitchInfo)
-        assert switch is not None, "invert_switch not found"
+        switch = require_entity(entities, "invert_switch", SwitchInfo)
 
         client.switch_command(switch.key, True)
         # The wire byte carries the inverted value...
@@ -1202,8 +1194,7 @@ async def test_uart_mock_modbus_deprecated_write_buffer(
         api_client_connected() as client,
     ):
         entities = await tracker.setup_and_start_scenario(client)
-        number = find_entity(entities, "buf_number", NumberInfo)
-        assert number is not None, "buf_number not found"
+        number = require_entity(entities, "buf_number", NumberInfo)
 
         # First write via the deprecated buffer path.
         client.number_command(number.key, 111)

tests/integration/test_uart_mock_pylontech.py

@@ -16,7 +16,7 @@ import asyncio
 from aioesphomeapi import ButtonInfo
 import pytest
 
-from .state_utils import InitialStateHelper, SensorStateCollector, find_entity
+from .state_utils import InitialStateHelper, SensorStateCollector, require_entity
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 SENSORS = [
@@ -102,8 +102,7 @@ async def test_uart_mock_pylontech(
         except TimeoutError:
             pytest.fail("Timeout waiting for initial states")
 
-        start_btn = find_entity(entities, "start_scenario", ButtonInfo)
-        assert start_btn is not None, "Start Scenario button not found"
+        start_btn = require_entity(entities, "start_scenario", ButtonInfo)
         client.button_command(start_btn.key)
 
         try: