Merge remote-tracking branch 'origin/dev' into assign-bits-helper

Co-authored-by: Jonathan Swoboda <154711427+swoboda1337@users.noreply.github.com>

AGENTS.md

@@ -211,8 +211,8 @@ file does, and it is the authority when they disagree. The most useful starting
         logged with `get_log_str` and `LOG_STR_ARG`, copied out with `ESPHOME_strncpy_P` before `publish_state`.
         A `dump_config` helper that picks a name returns `const LogString *` from `LOG_STR(...)`.
     *   **CRCs and bit tricks** come from `esphome/core/helpers.h` (`crc8`, `crc16`, `xor_bytes`, `sum16_bytes`,
-        `reverse_bits`, `encode_uint16`, `encode_value_le`, `sign_extend<BITS>`, `bcd_to_uint8`); pick the helper's
-        arguments instead of writing a loop. Epson style 7-byte RTC calendars (one-hot weekday) use
+        `reverse_bits`, `encode_uint16`, `encode_value_le`, `sign_extend<BITS>`, `bcd_to_uint8`, `assign_bits`); pick
+        the helper's arguments instead of writing a loop. Epson style 7-byte RTC calendars (one-hot weekday) use
         `time::decode_bcd_calendar`.
     *   **Bus helpers:** UART drains with `read_available(on_byte)` and `discard_available()`, text lines with
         `LineReader<N>`; I2C has `update_register`, `write_registers` (PROGMEM `I2CRegValue` table), `check_chip_id`

esphome/core/helpers.h

@@ -1067,6 +1067,16 @@ constexpr uint8_t bcd_to_uint8(uint8_t bcd) { return (bcd >> 4) * 10 + (bcd & 0x
 /// Encode a value from 0 to 99 as a two-digit packed BCD byte, e.g. 59 to 0x59.
 constexpr uint8_t uint8_to_bcd(uint8_t value) { return ((value / 10) << 4) | (value % 10); }
 
+/// Set the bits of \p mask in \p value when \p on is true, clear them otherwise.
+/// Use an unsigned mask for a register wider than int; a negative int mask is sign-extended.
+template<typename T> ESPHOME_ALWAYS_INLINE constexpr void assign_bits(T &value, std::type_identity_t<T> mask, bool on) {
+  if (on) {
+    value |= mask;
+  } else {
+    value &= ~mask;
+  }
+}
+
 /// Reverse the order of 8 bits.
 inline uint8_t reverse_bits(uint8_t x) {
   x = ((x & 0xAA) >> 1) | ((x & 0x55) << 1);

tests/components/core/test_assign_bits.cpp

@@ -0,0 +1,30 @@
+#include <gtest/gtest.h>
+
+#include "esphome/core/helpers.h"
+
+namespace esphome::core::testing {
+
+static_assert([] {
+  uint8_t v = 0;
+  assign_bits(v, 1 << 3, true);
+  return v;
+}() == 0x08);
+
+TEST(AssignBits, SetsAndClearsOnlyTheMask) {
+  uint8_t reg = 0b1010'0001;
+  assign_bits(reg, 0b0000'0110, true);
+  EXPECT_EQ(reg, 0b1010'0111);
+  assign_bits(reg, 0b1000'0010, false);
+  EXPECT_EQ(reg, 0b0010'0101);
+}
+
+TEST(AssignBits, ClearsAtTheTargetWidth) {
+  uint16_t reg16 = 0xFFFF;
+  assign_bits(reg16, 1 << 15, false);
+  EXPECT_EQ(reg16, 0x7FFF);
+  uint64_t reg64 = UINT64_MAX;
+  assign_bits(reg64, uint32_t{1} << 31, false);
+  EXPECT_EQ(reg64, UINT64_MAX & ~(uint64_t{1} << 31));
+}
+
+}  // namespace esphome::core::testing

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_crc8_helper.py

@@ -8,6 +8,7 @@ import pytest
 
 from .state_utils import assert_device_name
 from .types import APIClientConnectedFactory, RunCompiledFunction
+from .wait_utils import wait_or_fail
 
 
 @pytest.mark.asyncio
@@ -69,10 +70,9 @@ async def test_crc8_helper(
         await assert_device_name(client, "crc8-helper-test")
 
         # Wait for tests to complete with timeout
-        try:
-            await asyncio.wait_for(test_complete.wait(), timeout=5.0)
-        except TimeoutError:
-            pytest.fail("CRC8 integration test timed out after 5 seconds")
+        await wait_or_fail(
+            test_complete, 5.0, "CRC8 integration test timed out after 5 seconds"
+        )
 
         # Verify all tests passed
         assert test_results["setup_started"], "CRC8 test setup never started"

tests/integration/test_gpio_expander_cache.py

@@ -9,6 +9,7 @@ import pytest
 
 from .state_utils import assert_device_name
 from .types import APIClientConnectedFactory, RunCompiledFunction
+from .wait_utils import wait_or_fail
 
 
 @pytest.mark.asyncio
@@ -124,10 +125,7 @@ async def test_gpio_expander_cache(
         # Verify device info
         await assert_device_name(client, "gpio-expander-cache")
 
-        try:
-            await asyncio.wait_for(logs_done.wait(), timeout=5.0)
-        except TimeoutError:
-            pytest.fail("Timeout waiting for logs to complete")
+        await wait_or_fail(logs_done, 5.0, "Timeout waiting for logs to complete")
 
         assert index == len(log_order), (
             f"Expected {len(log_order)} log entries, but got {index}"

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_ota.py

@@ -13,7 +13,9 @@ from collections.abc import Generator
 from contextlib import contextmanager
 from dataclasses import dataclass
 import functools
+import logging
 from pathlib import Path
+import shutil
 import socket
 import zlib
 
@@ -34,6 +36,8 @@ from .host_prefs import force_safe_mode
 from .state_utils import assert_device_name
 from .types import APIClientConnectedFactory, CompileFunction, ConfigWriter
 
+_LOGGER = logging.getLogger(__name__)
+
 DEVICE_NAME = "host-ota-test"
 API_KEY = "AAECAwQFBgcICQoLDA0ODxAREhMUFRYXGBkaGxwdHh8="
 
@@ -80,6 +84,10 @@ def _udp_mode(
             real_ack(self, data)
 
     with monkeypatch.context() as m:
+        # Loopback answers in about a millisecond; resend lost datagrams
+        # sooner so the loss runs do not dominate the test time
+        m.setattr(espota2, "UDP_RESEND_INITIAL", 0.05)
+        m.setattr(espota2, "UDP_RESEND_MIN", 0.02)
         if force:
             m.setattr(espota2, "UDP_SLOW_EXCHANGE", -1.0)
         if blocked:
@@ -142,7 +150,23 @@ async def _build(
         binary_path = await compile_esphome(config_path)
         api_socket.close()
         ota_socket.close()
-    return api_port, ota_port, binary_path
+    return api_port, ota_port, await _strip_debug(binary_path)
+
+
+async def _strip_debug(binary_path: Path) -> Path:
+    """The fixtures build with -g; every upload sends the image, so drop the
+    debug info (GNU and Apple strip both take -S)."""
+    if (strip := shutil.which("strip")) is None:
+        _LOGGER.warning(
+            "strip not found; uploading the unstripped binary, which is much slower"
+        )
+        return binary_path
+    stripped = binary_path.with_name(f"{binary_path.name}.stripped")
+    proc = await asyncio.create_subprocess_exec(
+        strip, "-S", "-o", str(stripped), str(binary_path)
+    )
+    assert await proc.wait() == 0, "strip failed"
+    return stripped
 
 
 async def _run_ota(

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,8 +19,9 @@ 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
+from .wait_utils import wait_or_fail
 
 
 @pytest.mark.asyncio
@@ -61,9 +62,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)
@@ -90,10 +90,11 @@ async def test_light_constant_brightness(
             """Send a light command and wait for the state response."""
             state_futures[light_key] = loop.create_future()
             client.light_command(key=light_key, **kwargs)
-            try:
-                return await asyncio.wait_for(state_futures[light_key], timeout=timeout)
-            except TimeoutError:
-                pytest.fail(f"Timeout waiting for light state after command: {kwargs}")
+            return await wait_or_fail(
+                state_futures[light_key],
+                timeout,
+                f"Timeout waiting for light state after command: {kwargs}",
+            )
 
         # --- Test constant_brightness: true ---
 

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_loop_disable_enable.py

@@ -9,6 +9,7 @@ import pytest
 
 from .state_utils import assert_device_name
 from .types import APIClientConnectedFactory, RunCompiledFunction
+from .wait_utils import wait_or_fail
 
 
 @pytest.mark.asyncio
@@ -136,10 +137,11 @@ async def test_loop_disable_enable(
         await assert_device_name(client, "loop-test")
 
         # Wait for self_disable_10 to disable itself
-        try:
-            await asyncio.wait_for(self_disable_10_disabled.wait(), timeout=10.0)
-        except TimeoutError:
-            pytest.fail("self_disable_10 did not disable itself within 10 seconds")
+        await wait_or_fail(
+            self_disable_10_disabled,
+            10.0,
+            "self_disable_10 did not disable itself within 10 seconds",
+        )
 
         # Verify it ran at least 10 times before disabling
         assert len([c for c in self_disable_10_counts if c <= 10]) == 10, (
@@ -150,31 +152,32 @@ async def test_loop_disable_enable(
         )
 
         # Wait for normal_component to run at least 10 times
-        try:
-            await asyncio.wait_for(normal_component_10_loops.wait(), timeout=10.0)
-        except TimeoutError:
-            pytest.fail(
-                f"normal_component did not reach 10 loops within timeout, got {len(normal_component_counts)}"
-            )
+        await wait_or_fail(
+            normal_component_10_loops,
+            10.0,
+            "normal_component did not reach 10 loops within timeout",
+            normal_component_counts=normal_component_counts,
+        )
 
         # Wait for redundant operation tests
-        try:
-            await asyncio.wait_for(redundant_enable_tested.wait(), timeout=10.0)
-        except TimeoutError:
-            pytest.fail("redundant_enable did not test enabling when already enabled")
+        await wait_or_fail(
+            redundant_enable_tested,
+            10.0,
+            "redundant_enable did not test enabling when already enabled",
+        )
 
-        try:
-            await asyncio.wait_for(redundant_disable_tested.wait(), timeout=10.0)
-        except TimeoutError:
-            pytest.fail(
-                "redundant_disable did not test disabling when will be disabled"
-            )
+        await wait_or_fail(
+            redundant_disable_tested,
+            10.0,
+            "redundant_disable did not test disabling when will be disabled",
+        )
 
         # Wait to see if self_disable_10 gets re-enabled
-        try:
-            await asyncio.wait_for(self_disable_10_re_enabled.wait(), timeout=5.0)
-        except TimeoutError:
-            pytest.fail("self_disable_10 was not re-enabled within 5 seconds")
+        await wait_or_fail(
+            self_disable_10_re_enabled,
+            5.0,
+            "self_disable_10 was not re-enabled within 5 seconds",
+        )
 
         # Component was re-enabled - verify it ran more times
         later_self_disable_counts = [c for c in self_disable_10_counts if c > 10]
@@ -184,10 +187,11 @@ async def test_loop_disable_enable(
 
         # Test ISR component functionality
         # Wait for ISR component to disable itself after 5 loops
-        try:
-            await asyncio.wait_for(isr_component_disabled.wait(), timeout=3.0)
-        except TimeoutError:
-            pytest.fail("ISR component did not disable itself within 3 seconds")
+        await wait_or_fail(
+            isr_component_disabled,
+            3.0,
+            "ISR component did not disable itself within 3 seconds",
+        )
 
         # Verify it ran exactly 5 times before disabling
         first_run_counts = [c for c in isr_component_counts if c <= 5]
@@ -196,10 +200,11 @@ async def test_loop_disable_enable(
         )
 
         # Wait for component to be re-enabled by periodic ISR simulation and run again
-        try:
-            await asyncio.wait_for(isr_component_re_enabled.wait(), timeout=2.0)
-        except TimeoutError:
-            pytest.fail("ISR component was not re-enabled after ISR call")
+        await wait_or_fail(
+            isr_component_re_enabled,
+            2.0,
+            "ISR component was not re-enabled after ISR call",
+        )
 
         # Verify it's running again after ISR enable
         count_after_isr = len(isr_component_counts)
@@ -208,10 +213,11 @@ async def test_loop_disable_enable(
         )
 
         # Wait for pure ISR enable (no main loop enable) to work
-        try:
-            await asyncio.wait_for(isr_component_pure_re_enabled.wait(), timeout=2.0)
-        except TimeoutError:
-            pytest.fail("ISR component was not re-enabled by pure ISR call")
+        await wait_or_fail(
+            isr_component_pure_re_enabled,
+            2.0,
+            "ISR component was not re-enabled by pure ISR call",
+        )
 
         # Verify it ran after pure ISR enable
         final_count = len(isr_component_counts)
@@ -221,10 +227,11 @@ async def test_loop_disable_enable(
 
         # Test component.update functionality when loop is disabled
         # Wait for update component to disable its loop
-        try:
-            await asyncio.wait_for(update_component_loop_disabled.wait(), timeout=3.0)
-        except TimeoutError:
-            pytest.fail("Update component did not disable its loop within 3 seconds")
+        await wait_or_fail(
+            update_component_loop_disabled,
+            3.0,
+            "Update component did not disable its loop within 3 seconds",
+        )
 
         # Verify it ran exactly 3 loops before disabling
         assert update_component_loop_count == 3, (
@@ -232,12 +239,11 @@ async def test_loop_disable_enable(
         )
 
         # Wait for manual component.update to be called
-        try:
-            await asyncio.wait_for(
-                update_component_manual_update_called.wait(), timeout=5.0
-            )
-        except TimeoutError:
-            pytest.fail("Manual component.update was not called within 5 seconds")
+        await wait_or_fail(
+            update_component_manual_update_called,
+            5.0,
+            "Manual component.update was not called within 5 seconds",
+        )
 
         # The key test: verify that manual component.update worked after loop was disabled
         assert update_component_manual_update_count >= 1, (

tests/integration/test_micros_to_millis.py

@@ -9,6 +9,7 @@ import pytest
 
 from .state_utils import assert_device_name
 from .types import APIClientConnectedFactory, RunCompiledFunction
+from .wait_utils import wait_or_fail
 
 
 @pytest.mark.asyncio
@@ -35,11 +36,8 @@ async def test_micros_to_millis(
     ):
         await assert_device_name(client, "micros-to-millis-test")
 
-        try:
-            await asyncio.wait_for(all_passed.wait(), timeout=2.0)
-        except TimeoutError:
-            if failures:
-                pytest.fail(f"micros_to_millis failures: {failures}")
-            pytest.fail("micros_to_millis test timed out")
+        await wait_or_fail(
+            all_passed, 2.0, "micros_to_millis test timed out", failures=failures
+        )
 
         assert not failures, f"micros_to_millis failures: {failures}"

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_scheduler_background_wake.py

@@ -14,6 +14,7 @@ import pytest
 
 from .state_utils import assert_device_name, require_entity
 from .types import APIClientConnectedFactory, RunCompiledFunction
+from .wait_utils import wait_or_fail
 
 
 @pytest.mark.asyncio
@@ -57,12 +58,11 @@ async def test_scheduler_background_wake(
             )
             client.button_command(start_button.key)
 
-            try:
-                elapsed, loop_delta = await asyncio.wait_for(
-                    results[delay], timeout=10.0
-                )
-            except TimeoutError:
-                pytest.fail(f"background insert with delay={delay} did not fire")
+            elapsed, loop_delta = await wait_or_fail(
+                results[delay],
+                10.0,
+                f"background insert with delay={delay} did not fire",
+            )
 
             assert elapsed < 1000, (
                 f"background insert with delay={delay} should interrupt the "

tests/integration/test_scheduler_bulk_cleanup.py

@@ -7,6 +7,7 @@ import pytest
 
 from .state_utils import assert_device_name, require_service
 from .types import APIClientConnectedFactory, RunCompiledFunction
+from .wait_utils import wait_or_fail
 
 
 @pytest.mark.asyncio
@@ -80,10 +81,7 @@ async def test_scheduler_bulk_cleanup(
         await client.execute_service(trigger_bulk_cleanup_service, {})
 
         # Wait for test completion
-        try:
-            await asyncio.wait_for(test_complete_future, timeout=10.0)
-        except TimeoutError:
-            pytest.fail("Bulk cleanup test timed out")
+        await wait_or_fail(test_complete_future, 10.0, "Bulk cleanup test timed out")
 
         # Verify bulk cleanup was triggered
         assert bulk_cleanup_triggered, (

tests/integration/test_scheduler_defer_cancel.py

@@ -7,6 +7,7 @@ import pytest
 
 from .state_utils import assert_device_name, require_entity, require_service
 from .types import APIClientConnectedFactory, RunCompiledFunction
+from .wait_utils import wait_or_fail
 
 
 @pytest.mark.asyncio
@@ -70,11 +71,12 @@ async def test_scheduler_defer_cancel(
         await client.execute_service(test_defer_cancel_service, {})
 
         # Wait for test completion
-        try:
-            await asyncio.wait_for(test_complete_future, timeout=10.0)
-            executed_defer = await asyncio.wait_for(test_result_future, timeout=1.0)
-        except TimeoutError:
-            pytest.fail("Test did not complete within timeout")
+        await wait_or_fail(
+            test_complete_future, 10.0, "Test did not complete within timeout"
+        )
+        executed_defer = await wait_or_fail(
+            test_result_future, 1.0, "Test did not complete within timeout"
+        )
 
         # Verify that only defer 10 was executed
         assert executed_defer == 10, (

tests/integration/test_scheduler_defer_cancel_regular.py

@@ -6,6 +6,7 @@ import pytest
 
 from .state_utils import assert_device_name, require_service
 from .types import APIClientConnectedFactory, RunCompiledFunction
+from .wait_utils import wait_or_fail
 
 
 @pytest.mark.asyncio
@@ -54,10 +55,9 @@ async def test_scheduler_defer_cancels_regular(
         await client.execute_service(test_service, {})
 
         # Wait for test completion
-        try:
-            await asyncio.wait_for(test_complete_future, timeout=5.0)
-        except TimeoutError:
-            pytest.fail(f"Test timed out. Log messages: {log_messages}")
+        await wait_or_fail(
+            test_complete_future, 5.0, "Test timed out", log_messages=log_messages
+        )
 
         # Verify results
         assert not error_detected, (

tests/integration/test_scheduler_defer_fifo_simple.py

@@ -7,6 +7,7 @@ import pytest
 
 from .state_utils import assert_device_name, require_entity, require_service
 from .types import APIClientConnectedFactory, RunCompiledFunction
+from .wait_utils import wait_or_fail
 
 
 @pytest.mark.asyncio
@@ -70,11 +71,16 @@ async def test_scheduler_defer_fifo_simple(
         await client.execute_service(test_set_timeout_service, {})
 
         # Wait for first test completion
-        try:
-            await asyncio.wait_for(test_complete_future, timeout=5.0)
-            test1_passed = await asyncio.wait_for(test_result_future, timeout=1.0)
-        except TimeoutError:
-            pytest.fail("Test set_timeout(0) did not complete within 5 seconds")
+        await wait_or_fail(
+            test_complete_future,
+            5.0,
+            "Test set_timeout(0) did not complete within 5 seconds",
+        )
+        test1_passed = await wait_or_fail(
+            test_result_future,
+            1.0,
+            "Test set_timeout(0) did not complete within 5 seconds",
+        )
 
         assert test1_passed is True, (
             "set_timeout(0) FIFO test failed - items executed out of order"
@@ -88,11 +94,12 @@ async def test_scheduler_defer_fifo_simple(
         await client.execute_service(test_defer_service, {})
 
         # Wait for second test completion
-        try:
-            await asyncio.wait_for(test_complete_future, timeout=5.0)
-            test2_passed = await asyncio.wait_for(test_result_future, timeout=1.0)
-        except TimeoutError:
-            pytest.fail("Test defer() did not complete within 5 seconds")
+        await wait_or_fail(
+            test_complete_future, 5.0, "Test defer() did not complete within 5 seconds"
+        )
+        test2_passed = await wait_or_fail(
+            test_result_future, 1.0, "Test defer() did not complete within 5 seconds"
+        )
 
         # Verify the test passed
         assert test2_passed is True, (

tests/integration/test_scheduler_internal_id_no_collision.py

@@ -13,6 +13,7 @@ import pytest
 
 from .state_utils import assert_device_name
 from .types import APIClientConnectedFactory, RunCompiledFunction
+from .wait_utils import wait_or_fail
 
 
 @pytest.mark.asyncio
@@ -77,10 +78,11 @@ async def test_scheduler_internal_id_no_collision(
     ):
         await assert_device_name(client, "scheduler-internal-id-test")
 
-        try:
-            await asyncio.wait_for(all_tests_complete.wait(), timeout=5.0)
-        except TimeoutError:
-            pytest.fail("Not all collision tests completed within 5 seconds")
+        await wait_or_fail(
+            all_tests_complete,
+            5.0,
+            "Not all collision tests completed within 5 seconds",
+        )
 
         # Test 1: Both timeout types with same ID 0 must fire
         assert internal_timeout_0_fired.is_set(), (

tests/integration/test_scheduler_interval_reschedule.py

@@ -21,6 +21,7 @@ import pytest
 
 from .state_utils import assert_device_name
 from .types import APIClientConnectedFactory, RunCompiledFunction
+from .wait_utils import wait_or_fail
 
 
 @pytest.mark.asyncio
@@ -85,17 +86,15 @@ async def test_scheduler_interval_reschedule(
         await assert_device_name(client, "sched-interval-resched")
 
         # 1. Fast interval (50ms) should reach 10 fires within ~600ms
-        try:
-            await asyncio.wait_for(fast_10_future, timeout=5.0)
-        except TimeoutError:
-            pytest.fail("Fast interval (50ms) did not fire 10 times")
+        await wait_or_fail(
+            fast_10_future, 5.0, "Fast interval (50ms) did not fire 10 times"
+        )
 
         # 2. Medium interval (100ms) should reach 5 fires
         #    At that point, fast_count should be roughly 2x medium_count
-        try:
-            result = await asyncio.wait_for(medium_5_future, timeout=5.0)
-        except TimeoutError:
-            pytest.fail("Medium interval (100ms) did not fire 5 times")
+        result = await wait_or_fail(
+            medium_5_future, 5.0, "Medium interval (100ms) did not fire 5 times"
+        )
 
         fast_at_medium_5 = result[0]
         # Fast runs at 50ms, medium at 100ms, so fast should be ~2x medium
@@ -107,10 +106,9 @@ async def test_scheduler_interval_reschedule(
 
         # 3. Slow interval (200ms) should reach 3 fires
         #    At that point, both fast and medium should have proportionally more fires
-        try:
-            result = await asyncio.wait_for(slow_3_future, timeout=5.0)
-        except TimeoutError:
-            pytest.fail("Slow interval (200ms) did not fire 3 times")
+        result = await wait_or_fail(
+            slow_3_future, 5.0, "Slow interval (200ms) did not fire 3 times"
+        )
 
         fast_at_slow_3, medium_at_slow_3 = result
         # At 600ms: fast ~12, medium ~6, slow 3
@@ -124,10 +122,9 @@ async def test_scheduler_interval_reschedule(
         )
 
         # 4. Self-cancelling interval should have stopped after exactly 3 fires
-        try:
-            await asyncio.wait_for(self_cancel_stopped_future, timeout=5.0)
-        except TimeoutError:
-            pytest.fail("Self-cancelling interval did not stop")
+        await wait_or_fail(
+            self_cancel_stopped_future, 5.0, "Self-cancelling interval did not stop"
+        )
 
         # Wait a bit to ensure it doesn't fire again
         await asyncio.sleep(0.3)
@@ -137,18 +134,18 @@ async def test_scheduler_interval_reschedule(
         )
 
         # 5. Timeout scheduled from interval callback should have fired
-        try:
-            await asyncio.wait_for(callback_timeout_future, timeout=5.0)
-        except TimeoutError:
-            pytest.fail("Timeout scheduled from interval callback did not fire")
+        await wait_or_fail(
+            callback_timeout_future,
+            5.0,
+            "Timeout scheduled from interval callback did not fire",
+        )
 
         # 6. Interval that replaces itself via set_interval from within callback
         #    The original fires once, sets up a new named interval, then stops itself.
         #    The replacement interval should fire 3 times then cancel itself.
-        try:
-            await asyncio.wait_for(replace_original_future, timeout=5.0)
-        except TimeoutError:
-            pytest.fail("Replace-test original interval did not fire")
+        await wait_or_fail(
+            replace_original_future, 5.0, "Replace-test original interval did not fire"
+        )
 
         try:
             await asyncio.wait_for(replaced_stopped_future, timeout=5.0)

tests/integration/test_scheduler_interval_zero_coerced.py

@@ -19,6 +19,7 @@ import pytest
 
 from .state_utils import assert_device_name
 from .types import APIClientConnectedFactory, RunCompiledFunction
+from .wait_utils import wait_or_fail
 
 
 @pytest.mark.asyncio
@@ -48,19 +49,19 @@ async def test_scheduler_interval_zero_coerced(
         await assert_device_name(client, "sched-interval-zero")
 
         # Coerce warning must fire at registration
-        try:
-            await asyncio.wait_for(coerce_warning, timeout=5.0)
-        except TimeoutError:
-            pytest.fail("Expected coerce warning 'would spin main loop' not seen")
+        await wait_or_fail(
+            coerce_warning,
+            5.0,
+            "Expected coerce warning 'would spin main loop' not seen",
+        )
 
         # The coerced 1ms interval should fire 50 times quickly — this
         # confirms the callback actually runs (not just registered) and the
         # scheduler yields back to the main loop each time.
-        try:
-            await asyncio.wait_for(reached_50, timeout=5.0)
-        except TimeoutError:
-            pytest.fail(
-                "Coerced interval=0→1ms did not reach 50 fires within 5s, "
-                "which would indicate either the coercion failed or the "
-                "main loop is still being starved."
-            )
+        await wait_or_fail(
+            reached_50,
+            5.0,
+            "Coerced interval=0→1ms did not reach 50 fires within 5s, "
+            "which would indicate either the coercion failed or the "
+            "main loop is still being starved.",
+        )

tests/integration/test_scheduler_loop_phase_timeout.py

@@ -9,6 +9,7 @@ import pytest
 
 from .conftest import run_binary
 from .types import CompileFunction, ConfigWriter
+from .wait_utils import wait_or_fail
 
 
 @pytest.mark.asyncio
@@ -29,10 +30,9 @@ async def test_scheduler_loop_phase_timeout(
             timeout_result.set_result(int(match.group(1)))
 
     async with run_binary(binary_path, line_callback=on_log_line):
-        try:
-            elapsed = await asyncio.wait_for(timeout_result, timeout=15.0)
-        except TimeoutError:
-            pytest.fail("loop-phase timeout did not fire within 15 seconds")
+        elapsed = await wait_or_fail(
+            timeout_result, 15.0, "loop-phase timeout did not fire within 15 seconds"
+        )
 
         assert elapsed < 1000, (
             f"loop-phase timeout should interrupt the five-second sleep, but fired after {elapsed}ms"

tests/integration/test_scheduler_null_name.py

@@ -7,6 +7,7 @@ import pytest
 
 from .state_utils import assert_device_name, require_service
 from .types import APIClientConnectedFactory, RunCompiledFunction
+from .wait_utils import wait_or_fail
 
 
 @pytest.mark.asyncio
@@ -47,9 +48,8 @@ async def test_scheduler_null_name(
         await client.execute_service(test_null_name_service, {})
 
         # Wait for test completion
-        try:
-            await asyncio.wait_for(test_complete_future, timeout=10.0)
-        except TimeoutError:
-            pytest.fail(
-                "Test did not complete within timeout - likely crashed due to NULL name"
-            )
+        await wait_or_fail(
+            test_complete_future,
+            10.0,
+            "Test did not complete within timeout - likely crashed due to NULL name",
+        )

tests/integration/test_scheduler_numeric_id_test.py

@@ -7,6 +7,7 @@ import pytest
 
 from .state_utils import assert_device_name
 from .types import APIClientConnectedFactory, RunCompiledFunction
+from .wait_utils import wait_or_fail
 
 
 @pytest.mark.asyncio
@@ -124,25 +125,29 @@ async def test_scheduler_numeric_id_test(
         await assert_device_name(client, "scheduler-numeric-id-test")
 
         # Wait for numeric timeout tests
-        try:
-            await asyncio.wait_for(numeric_timeout_1001_fired.wait(), timeout=0.5)
-        except TimeoutError:
-            pytest.fail("Numeric timeout 1001 did not fire within 0.5 seconds")
+        await wait_or_fail(
+            numeric_timeout_1001_fired,
+            0.5,
+            "Numeric timeout 1001 did not fire within 0.5 seconds",
+        )
 
-        try:
-            await asyncio.wait_for(numeric_timeout_1002_fired.wait(), timeout=0.5)
-        except TimeoutError:
-            pytest.fail("Numeric timeout 1002 did not fire within 0.5 seconds")
+        await wait_or_fail(
+            numeric_timeout_1002_fired,
+            0.5,
+            "Numeric timeout 1002 did not fire within 0.5 seconds",
+        )
 
-        try:
-            await asyncio.wait_for(numeric_interval_2001_fired.wait(), timeout=1.0)
-        except TimeoutError:
-            pytest.fail("Numeric interval 2001 did not fire within 1 second")
+        await wait_or_fail(
+            numeric_interval_2001_fired,
+            1.0,
+            "Numeric interval 2001 did not fire within 1 second",
+        )
 
-        try:
-            await asyncio.wait_for(numeric_interval_cancelled.wait(), timeout=2.0)
-        except TimeoutError:
-            pytest.fail("Numeric interval 2001 was not cancelled within 2 seconds")
+        await wait_or_fail(
+            numeric_interval_cancelled,
+            2.0,
+            "Numeric interval 2001 was not cancelled within 2 seconds",
+        )
 
         # Verify numeric interval ran at least twice
         assert numeric_interval_count >= 2, (
@@ -155,55 +160,60 @@ async def test_scheduler_numeric_id_test(
         )
 
         # Wait for duplicate timeout (only last one should fire)
-        try:
-            await asyncio.wait_for(duplicate_timeout_fired.wait(), timeout=1.0)
-        except TimeoutError:
-            pytest.fail("Duplicate numeric timeout did not fire within 1 second")
+        await wait_or_fail(
+            duplicate_timeout_fired,
+            1.0,
+            "Duplicate numeric timeout did not fire within 1 second",
+        )
 
         # Wait for component method tests
-        try:
-            await asyncio.wait_for(component_timeout_fired.wait(), timeout=0.5)
-        except TimeoutError:
-            pytest.fail("Component numeric timeout did not fire within 0.5 seconds")
+        await wait_or_fail(
+            component_timeout_fired,
+            0.5,
+            "Component numeric timeout did not fire within 0.5 seconds",
+        )
 
-        try:
-            await asyncio.wait_for(component_interval_fired.wait(), timeout=1.0)
-        except TimeoutError:
-            pytest.fail("Component numeric interval did not fire within 1 second")
+        await wait_or_fail(
+            component_interval_fired,
+            1.0,
+            "Component numeric interval did not fire within 1 second",
+        )
 
         # Wait for edge case tests
-        try:
-            await asyncio.wait_for(zero_id_timeout_fired.wait(), timeout=0.5)
-        except TimeoutError:
-            pytest.fail("Zero ID timeout did not fire within 0.5 seconds")
+        await wait_or_fail(
+            zero_id_timeout_fired,
+            0.5,
+            "Zero ID timeout did not fire within 0.5 seconds",
+        )
 
-        try:
-            await asyncio.wait_for(max_id_timeout_fired.wait(), timeout=0.5)
-        except TimeoutError:
-            pytest.fail("Max ID timeout did not fire within 0.5 seconds")
+        await wait_or_fail(
+            max_id_timeout_fired, 0.5, "Max ID timeout did not fire within 0.5 seconds"
+        )
 
         # Wait for numeric defer tests
-        try:
-            await asyncio.wait_for(numeric_defer_7001_fired.wait(), timeout=0.5)
-        except TimeoutError:
-            pytest.fail("Numeric defer 7001 did not fire within 0.5 seconds")
+        await wait_or_fail(
+            numeric_defer_7001_fired,
+            0.5,
+            "Numeric defer 7001 did not fire within 0.5 seconds",
+        )
 
-        try:
-            await asyncio.wait_for(numeric_defer_7002_fired.wait(), timeout=0.5)
-        except TimeoutError:
-            pytest.fail("Numeric defer 7002 did not fire within 0.5 seconds")
+        await wait_or_fail(
+            numeric_defer_7002_fired,
+            0.5,
+            "Numeric defer 7002 did not fire within 0.5 seconds",
+        )
 
         # Verify numeric defer was cancelled
-        try:
-            await asyncio.wait_for(numeric_defer_cancelled.wait(), timeout=0.5)
-        except TimeoutError:
-            pytest.fail("Numeric defer 8001 cancel confirmation not received")
+        await wait_or_fail(
+            numeric_defer_cancelled,
+            0.5,
+            "Numeric defer 8001 cancel confirmation not received",
+        )
 
         # Wait for final results
-        try:
-            await asyncio.wait_for(final_results_logged.wait(), timeout=3.0)
-        except TimeoutError:
-            pytest.fail("Final results were not logged within 3 seconds")
+        await wait_or_fail(
+            final_results_logged, 3.0, "Final results were not logged within 3 seconds"
+        )
 
         # Verify results
         assert timeout_count >= 6, f"Expected at least 6 timeouts, got {timeout_count}"

tests/integration/test_scheduler_rapid_cancellation.py

@@ -7,6 +7,7 @@ import pytest
 
 from .state_utils import assert_device_name, require_service
 from .types import APIClientConnectedFactory, RunCompiledFunction
+from .wait_utils import wait_or_fail
 
 
 @pytest.mark.asyncio
@@ -90,10 +91,9 @@ async def test_scheduler_rapid_cancellation(
         await client.execute_service(run_test_service, {})
 
         # Wait for test to complete with timeout
-        try:
-            await asyncio.wait_for(test_complete_future, timeout=10.0)
-        except TimeoutError:
-            pytest.fail(f"Test timed out. Stats: {test_stats}")
+        await wait_or_fail(
+            test_complete_future, 10.0, "Test timed out", test_stats=test_stats
+        )
 
         # Check for any errors
         assert len(test_stats["errors"]) == 0, (

tests/integration/test_scheduler_recursive_timeout.py

@@ -6,6 +6,7 @@ import pytest
 
 from .state_utils import assert_device_name, require_service
 from .types import APIClientConnectedFactory, RunCompiledFunction
+from .wait_utils import wait_or_fail
 
 
 @pytest.mark.asyncio
@@ -61,12 +62,12 @@ async def test_scheduler_recursive_timeout(
         await client.execute_service(run_test_service, {})
 
         # Wait for test to complete
-        try:
-            await asyncio.wait_for(test_complete_future, timeout=10.0)
-        except TimeoutError:
-            pytest.fail(
-                f"Recursive timeout test timed out. Got sequence: {execution_sequence}"
-            )
+        await wait_or_fail(
+            test_complete_future,
+            10.0,
+            "Recursive timeout test timed out",
+            execution_sequence=execution_sequence,
+        )
 
         # Verify execution sequence
         assert execution_sequence == expected_sequence, (

tests/integration/test_scheduler_removed_item_race.py

@@ -7,6 +7,7 @@ import pytest
 
 from .state_utils import assert_device_name, require_service
 from .types import APIClientConnectedFactory, RunCompiledFunction
+from .wait_utils import wait_or_fail
 
 
 @pytest.mark.asyncio
@@ -82,10 +83,9 @@ async def test_scheduler_removed_item_race(
         await client.execute_service(run_test_service, {})
 
         # Wait for test completion
-        try:
-            await asyncio.wait_for(test_complete_future, timeout=5.0)
-        except TimeoutError:
-            pytest.fail("Test did not complete within timeout")
+        await wait_or_fail(
+            test_complete_future, 5.0, "Test did not complete within timeout"
+        )
 
         # Verify results
         assert test_passed, (

tests/integration/test_scheduler_self_keyed.py

@@ -14,6 +14,7 @@ import pytest
 
 from .state_utils import assert_device_name
 from .types import APIClientConnectedFactory, RunCompiledFunction
+from .wait_utils import wait_or_fail
 
 
 @pytest.mark.asyncio
@@ -64,10 +65,11 @@ async def test_scheduler_self_keyed(
     ):
         await assert_device_name(client, "scheduler-self-keyed-test")
 
-        try:
-            await asyncio.wait_for(all_tests_complete.wait(), timeout=5.0)
-        except TimeoutError:
-            pytest.fail("Not all self-keyed tests completed within 5 seconds")
+        await wait_or_fail(
+            all_tests_complete,
+            5.0,
+            "Not all self-keyed tests completed within 5 seconds",
+        )
 
         # Test 1: self-keyed timeout fires
         assert self_a_fired.is_set(), "Self timeout A should have fired"

tests/integration/test_scheduler_simultaneous_callbacks.py

@@ -7,6 +7,7 @@ import pytest
 
 from .state_utils import assert_device_name, require_service
 from .types import APIClientConnectedFactory, RunCompiledFunction
+from .wait_utils import wait_or_fail
 
 
 @pytest.mark.asyncio
@@ -80,10 +81,12 @@ async def test_scheduler_simultaneous_callbacks(
         await client.execute_service(run_test_service, {})
 
         # Wait for test to complete
-        try:
-            await asyncio.wait_for(test_complete_future, timeout=30.0)
-        except TimeoutError:
-            pytest.fail(f"Simultaneous callbacks test timed out. Stats: {test_stats}")
+        await wait_or_fail(
+            test_complete_future,
+            30.0,
+            "Simultaneous callbacks test timed out",
+            test_stats=test_stats,
+        )
 
         # Check for any errors
         assert len(test_stats["errors"]) == 0, (

tests/integration/test_scheduler_string_test.py

@@ -7,6 +7,7 @@ import pytest
 
 from .state_utils import assert_device_name
 from .types import APIClientConnectedFactory, RunCompiledFunction
+from .wait_utils import wait_or_fail
 
 
 @pytest.mark.asyncio
@@ -119,25 +120,27 @@ async def test_scheduler_string_test(
         await assert_device_name(client, "scheduler-string-test")
 
         # Wait for static string tests
-        try:
-            await asyncio.wait_for(static_timeout_1_fired.wait(), timeout=0.5)
-        except TimeoutError:
-            pytest.fail("Static timeout 1 did not fire within 0.5 seconds")
+        await wait_or_fail(
+            static_timeout_1_fired,
+            0.5,
+            "Static timeout 1 did not fire within 0.5 seconds",
+        )
 
-        try:
-            await asyncio.wait_for(static_timeout_2_fired.wait(), timeout=0.5)
-        except TimeoutError:
-            pytest.fail("Static timeout 2 did not fire within 0.5 seconds")
+        await wait_or_fail(
+            static_timeout_2_fired,
+            0.5,
+            "Static timeout 2 did not fire within 0.5 seconds",
+        )
 
-        try:
-            await asyncio.wait_for(static_interval_fired.wait(), timeout=1.0)
-        except TimeoutError:
-            pytest.fail("Static interval did not fire within 1 second")
+        await wait_or_fail(
+            static_interval_fired, 1.0, "Static interval did not fire within 1 second"
+        )
 
-        try:
-            await asyncio.wait_for(static_interval_cancelled.wait(), timeout=2.0)
-        except TimeoutError:
-            pytest.fail("Static interval was not cancelled within 2 seconds")
+        await wait_or_fail(
+            static_interval_cancelled,
+            2.0,
+            "Static interval was not cancelled within 2 seconds",
+        )
 
         # Verify static interval ran at least 3 times
         assert static_interval_count >= 2, (
@@ -150,44 +153,39 @@ async def test_scheduler_string_test(
         )
 
         # Wait for static defer tests
-        try:
-            await asyncio.wait_for(static_defer_1_fired.wait(), timeout=0.5)
-        except TimeoutError:
-            pytest.fail("Static defer 1 did not fire within 0.5 seconds")
+        await wait_or_fail(
+            static_defer_1_fired, 0.5, "Static defer 1 did not fire within 0.5 seconds"
+        )
 
-        try:
-            await asyncio.wait_for(static_defer_2_fired.wait(), timeout=0.5)
-        except TimeoutError:
-            pytest.fail("Static defer 2 did not fire within 0.5 seconds")
+        await wait_or_fail(
+            static_defer_2_fired, 0.5, "Static defer 2 did not fire within 0.5 seconds"
+        )
 
         # Wait for dynamic string tests
-        try:
-            await asyncio.wait_for(dynamic_timeout_fired.wait(), timeout=1.0)
-        except TimeoutError:
-            pytest.fail("Dynamic timeout did not fire within 1 second")
+        await wait_or_fail(
+            dynamic_timeout_fired, 1.0, "Dynamic timeout did not fire within 1 second"
+        )
 
-        try:
-            await asyncio.wait_for(dynamic_interval_fired.wait(), timeout=1.5)
-        except TimeoutError:
-            pytest.fail("Dynamic interval did not fire within 1.5 seconds")
+        await wait_or_fail(
+            dynamic_interval_fired,
+            1.5,
+            "Dynamic interval did not fire within 1.5 seconds",
+        )
 
         # Wait for dynamic defer test
-        try:
-            await asyncio.wait_for(dynamic_defer_fired.wait(), timeout=1.0)
-        except TimeoutError:
-            pytest.fail("Dynamic defer did not fire within 1 second")
+        await wait_or_fail(
+            dynamic_defer_fired, 1.0, "Dynamic defer did not fire within 1 second"
+        )
 
         # Wait for cancel test
-        try:
-            await asyncio.wait_for(cancel_test_done.wait(), timeout=1.0)
-        except TimeoutError:
-            pytest.fail("Cancel test did not complete within 1 second")
+        await wait_or_fail(
+            cancel_test_done, 1.0, "Cancel test did not complete within 1 second"
+        )
 
         # Wait for final results
-        try:
-            await asyncio.wait_for(final_results_logged.wait(), timeout=4.0)
-        except TimeoutError:
-            pytest.fail("Final results were not logged within 4 seconds")
+        await wait_or_fail(
+            final_results_logged, 4.0, "Final results were not logged within 4 seconds"
+        )
 
         # Verify results
         assert timeout_count >= 6, (

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_socket_tcp_client_link.py

@@ -85,11 +85,9 @@ async def test_socket_tcp_client_link(
             assert echoed and echoed[0] == PAYLOAD, "Echo payload mismatch"
 
             await wait_or_fail(link_down, 15.0, "Link never reported the drop")
-            try:
-                await asyncio.wait_for(reconnected.wait(), timeout=15.0)
-                await asyncio.wait_for(second_link_up.wait(), timeout=15.0)
-            except TimeoutError:
-                pytest.fail("Link did not reconnect after the server dropped it")
+            msg = "Link did not reconnect after the server dropped it"
+            await wait_or_fail(reconnected, 15.0, msg)
+            await wait_or_fail(second_link_up, 15.0, msg)
             assert second_echoed == [SECOND_PAYLOAD], (
                 "Second session echo wrong; stale bytes from the first session?"
             )

tests/integration/test_syslog.py

@@ -15,6 +15,7 @@ import pytest
 
 from .state_utils import assert_device_name, require_service
 from .types import APIClientConnectedFactory, RunCompiledFunction
+from .wait_utils import wait_or_fail
 
 
 class ParsedSyslogMessage(TypedDict):
@@ -69,10 +70,6 @@ class SyslogReceiver:
             if pattern.search(msg):
                 event.set()
 
-    async def wait_for_messages(self, timeout: float = 10.0) -> None:
-        """Wait for at least one message to be received."""
-        await asyncio.wait_for(self.message_received.wait(), timeout=timeout)
-
     async def wait_for_pattern(self, pattern: str, timeout: float = 5.0) -> str:
         """Wait for a message matching the pattern."""
         compiled = re.compile(pattern)
@@ -176,10 +173,11 @@ async def test_syslog(
             await assert_device_name(client, "syslog-test")
 
             # Wait for syslog messages (ESPHome logs during startup)
-            try:
-                await receiver.wait_for_messages(timeout=10.0)
-            except TimeoutError:
-                pytest.fail("No syslog messages received within timeout")
+            await wait_or_fail(
+                receiver.message_received,
+                10.0,
+                "No syslog messages received within timeout",
+            )
 
             # Give it a moment to collect more messages
             await asyncio.sleep(0.5)

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_dfrobot_sen0609.py

@@ -13,6 +13,7 @@ import pytest
 
 from .state_utils import require_entity
 from .types import APIClientConnectedFactory, RunCompiledFunction
+from .wait_utils import wait_or_fail
 
 
 @pytest.mark.asyncio
@@ -63,10 +64,7 @@ async def test_uart_mock_dfrobot_sen0609(
         assert state_of(min_range) == pytest.approx(0.6)
 
         client.number_command(max_range.key, 8.0)
-        try:
-            await asyncio.wait_for(saved, timeout=10.0)
-        except TimeoutError:
-            pytest.fail("The new range was not written and saved")
+        await wait_or_fail(saved, 10.0, "The new range was not written and saved")
         await wait_for(
             lambda: state_of(max_range) == pytest.approx(8.0), "new maximum range"
         )

tests/integration/test_uart_mock_ld2410.py

@@ -23,8 +23,9 @@ 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
+from .wait_utils import wait_or_fail
 
 
 @pytest.mark.asyncio
@@ -95,8 +96,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
@@ -118,13 +118,12 @@ async def test_uart_mock_ld2410(
 
         # Wait for the recovery frame (Phase 5) to be parsed
         # This proves the component survived garbage + truncated + overflow
-        try:
-            await asyncio.wait_for(recovery_received, timeout=15.0)
-        except TimeoutError:
-            pytest.fail(
-                f"Timeout waiting for recovery frame. Received:\n"
-                f"  sensor_states: {collector.sensor_states}"
-            )
+        await wait_or_fail(
+            recovery_received,
+            15.0,
+            "Timeout waiting for recovery frame",
+            sensor_states=collector.sensor_states,
+        )
 
         # Verify overflow warning was logged
         assert overflow_seen.done(), (
@@ -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
@@ -245,12 +243,11 @@ async def test_uart_mock_ld2410_engineering(
         assert collector.binary_states["out_pin_presence"][0] is True
 
         # Wait for Phase 3 frame (still_distance = 291cm, multi-byte)
-        try:
-            await asyncio.wait_for(phase3_received, timeout=15.0)
-        except TimeoutError:
-            pytest.fail(
-                f"Timeout waiting for Phase 3 frame. Received:\n"
-                f"  still_distance: {collector.sensor_states['still_distance']}"
-            )
+        await wait_or_fail(
+            phase3_received,
+            15.0,
+            "Timeout waiting for Phase 3 frame",
+            still_distance=collector.sensor_states["still_distance"],
+        )
 
         assert pytest.approx(291.0) in collector.sensor_states["still_distance"]

tests/integration/test_uart_mock_ld2410_query_ack.py

@@ -20,8 +20,9 @@ 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
+from .wait_utils import wait_or_fail
 
 # Number object_id: (value from the seed ACK, value from the ACK under test)
 EXPECTED: dict[str, tuple[float, float]] = {
@@ -124,26 +125,22 @@ 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:
-            await asyncio.wait_for(data_received, timeout=5.0)
-        except TimeoutError:
-            pytest.fail(
-                f"Timeout waiting for the data frame values. Received:\n"
-                f"  sensor_states: {sensor_states}"
-            )
-
-        try:
-            await asyncio.wait_for(final_received, timeout=5.0)
-        except TimeoutError:
-            pytest.fail(
-                f"Timeout waiting for the query ACK values. Received:\n"
-                f"  number_states: {number_states}\n"
-                f"  ack_errors: {ack_errors}"
-            )
+        await wait_or_fail(
+            data_received,
+            5.0,
+            "Timeout waiting for the data frame values",
+            sensor_states=sensor_states,
+        )
+        await wait_or_fail(
+            final_received,
+            5.0,
+            "Timeout waiting for the query ACK values",
+            number_states=number_states,
+            ack_errors=ack_errors,
+        )
 
         # One data frame, so each sensor is published once
         for name, value in SENSOR_EXPECTED.items():

tests/integration/test_uart_mock_ld2410s.py

@@ -43,6 +43,7 @@ import pytest
 
 from .state_utils import InitialStateHelper, SensorStateCollector, require_entity
 from .types import APIClientConnectedFactory, RunCompiledFunction
+from .wait_utils import wait_or_fail
 
 CONFIG_START_TX = "FD:FC:FB:FA:04:00:FF:00:01:00:04:03:02:01"
 VERSION_READ_TX = "FD:FC:FB:FA:02:00:00:00:04:03:02:01"
@@ -127,13 +128,12 @@ async def _start_scenario(client: APIClient, collector: SensorStateCollector) ->
 async def _wait_for_presence(
     collector: SensorStateCollector, waiter: asyncio.Future[bool]
 ) -> None:
-    try:
-        await asyncio.wait_for(waiter, timeout=5.0)
-    except TimeoutError:
-        pytest.fail(
-            f"Timeout waiting for presence updates. Received: "
-            f"{collector.binary_states['has_target']}"
-        )
+    await wait_or_fail(
+        waiter,
+        5.0,
+        "Timeout waiting for presence updates",
+        has_target=collector.binary_states["has_target"],
+    )
 
 
 @pytest.mark.asyncio
@@ -184,10 +184,12 @@ async def test_uart_mock_ld2410s(
         run_compiled(yaml_config, line_callback=line_callback),
         api_client_connected() as client,
     ):
-        try:
-            await asyncio.wait_for(asyncio.gather(setup_done, warning_cleared), 5.0)
-        except TimeoutError:
-            pytest.fail(f"Init sequence did not finish. TX frames: {tx_frames}")
+        await wait_or_fail(
+            asyncio.gather(setup_done, warning_cleared),
+            5.0,
+            "Init sequence did not finish",
+            tx_frames=tx_frames,
+        )
 
         assert tx_frames[:6] == INIT_TX
         assert parsed_answers == PARSED_ANSWERS
@@ -205,10 +207,9 @@ async def test_uart_mock_ld2410s(
         assert initial[no_delay_key].state == pytest.approx(10)
 
         client.number_command(max_distance_key, 7.0)
-        try:
-            await asyncio.wait_for(write_done, 5.0)
-        except TimeoutError:
-            pytest.fail(f"Parameter write not sent. TX frames: {tx_frames}")
+        await wait_or_fail(
+            write_done, 5.0, "Parameter write not sent", tx_frames=tx_frames
+        )
         assert tx_frames[6:9] == WRITE_TX, tx_frames
         assert number_states == [pytest.approx(7.0)]
 
@@ -279,10 +280,12 @@ async def test_uart_mock_ld2410s_recovery(
     ):
         await _start_scenario(client, collector)
         await _wait_for_presence(collector, all_presence)
-        try:
-            await asyncio.wait_for(asyncio.gather(resent, gave_up), 5.0)
-        except TimeoutError:
-            pytest.fail(f"No resend or give-up warning. TX frames: {tx_frames}")
+        await wait_or_fail(
+            asyncio.gather(resent, gave_up),
+            5.0,
+            "No resend or give-up warning",
+            tx_frames=tx_frames,
+        )
         # Let any unexpected extra update arrive before comparing
         await asyncio.sleep(0.3)
 
@@ -358,30 +361,28 @@ async def test_uart_mock_ld2410s_link_lost(
         run_compiled(yaml_config, line_callback=line_callback),
         api_client_connected() as client,
     ):
-        try:
-            await asyncio.wait_for(setup_done, 5.0)
-        except TimeoutError:
-            pytest.fail(f"Init sequence did not finish. TX frames: {tx_frames}")
+        await wait_or_fail(
+            setup_done, 5.0, "Init sequence did not finish", tx_frames=tx_frames
+        )
 
         await _start_scenario(client, collector)
-        try:
-            await asyncio.wait_for(all_presence, timeout=10.0)
-        except TimeoutError:
-            pytest.fail(
-                f"Timeout waiting for presence updates. Received: "
-                f"{collector.binary_states['has_target']}"
-            )
+        await wait_or_fail(
+            all_presence,
+            10.0,
+            "Timeout waiting for presence updates",
+            has_target=collector.binary_states["has_target"],
+        )
         assert collector.binary_states["has_target"] == EXPECTED_LINK_PRESENCE
 
-        try:
-            lost_at = await asyncio.wait_for(link_lost, timeout=7.0)
-            await asyncio.wait_for(
-                asyncio.gather(warning_set_after_loss, restarted), timeout=1.0
-            )
-        except TimeoutError:
-            pytest.fail(
-                f"No link-lost warning, warning flag or restart. TX frames: {tx_frames}"
-            )
+        lost_at = await wait_or_fail(
+            link_lost, 7.0, "No link-lost warning", tx_frames=tx_frames
+        )
+        await wait_or_fail(
+            asyncio.gather(warning_set_after_loss, restarted),
+            1.0,
+            "No warning flag or restart after the link was lost",
+            tx_frames=tx_frames,
+        )
 
         # The stream ran past the link timeout, so only the silence after it counts
         silence = lost_at - rx_injected_at[-1]

tests/integration/test_uart_mock_ld2412.py

@@ -46,10 +46,10 @@ from .state_utils import (
     InitialStateHelper,
     SensorStateCollector,
     StateWaiter,
-    find_entity,
     require_entity,
 )
 from .types import APIClientConnectedFactory, RunCompiledFunction
+from .wait_utils import wait_or_fail
 
 # Length and command word of each frame as the uart mock logs it; the payload is logged on its own TX line
 MOTION_COMMAND = "10:00:03:00"
@@ -193,8 +193,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
@@ -216,13 +215,12 @@ async def test_uart_mock_ld2412(
 
         # Wait for the recovery frame (Phase 5) to be parsed
         # This proves the component survived garbage + truncated + overflow
-        try:
-            await asyncio.wait_for(recovery_received, timeout=3.0)
-        except TimeoutError:
-            pytest.fail(
-                f"Timeout waiting for recovery frame. Received:\n"
-                f"  sensor_states: {collector.sensor_states}"
-            )
+        await wait_or_fail(
+            recovery_received,
+            3.0,
+            "Timeout waiting for recovery frame",
+            sensor_states=collector.sensor_states,
+        )
 
         # Verify overflow warning was logged
         assert overflow_seen.done(), (
@@ -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
@@ -324,24 +321,22 @@ async def test_uart_mock_ld2412_engineering(
         assert collector.sensor_states["light"][0] == pytest.approx(87.0)
 
         # Wait for Phase 3 frame: still_distance = 291cm (multi-byte)
-        try:
-            await asyncio.wait_for(phase3_still_received, timeout=3.0)
-        except TimeoutError:
-            pytest.fail(
-                f"Timeout waiting for Phase 3 still_distance. Received:\n"
-                f"  still_distance: {collector.sensor_states['still_distance']}"
-            )
+        await wait_or_fail(
+            phase3_still_received,
+            3.0,
+            "Timeout waiting for Phase 3 still_distance",
+            still_distance=collector.sensor_states["still_distance"],
+        )
 
         assert pytest.approx(291.0) in collector.sensor_states["still_distance"]
 
         # Wait for Phase 3: detection_distance = 291 (still-only target)
-        try:
-            await asyncio.wait_for(phase3_detect_received, timeout=3.0)
-        except TimeoutError:
-            pytest.fail(
-                f"Timeout waiting for detection_distance=291. "
-                f"Received: {collector.sensor_states['detection_distance']}"
-            )
+        await wait_or_fail(
+            phase3_detect_received,
+            3.0,
+            "Timeout waiting for detection_distance=291",
+            detection_distance=collector.sensor_states["detection_distance"],
+        )
 
         assert pytest.approx(291.0) in collector.sensor_states["detection_distance"]
 
@@ -407,8 +402,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
@@ -426,14 +420,13 @@ async def test_uart_mock_ld2412_engineering_truncated(
         assert collector.sensor_states["light"][0] == pytest.approx(87.0)
 
         # Wait for Phase 3 recovery frame (gate_0_move=50)
-        try:
-            await asyncio.wait_for(recovery_received, timeout=3.0)
-        except TimeoutError:
-            pytest.fail(
-                f"Timeout waiting for recovery frame. Received:\n"
-                f"  gate_0_move_energy: {collector.sensor_states['gate_0_move_energy']}\n"
-                f"  light: {collector.sensor_states['light']}"
-            )
+        await wait_or_fail(
+            recovery_received,
+            3.0,
+            "Timeout waiting for recovery frame",
+            gate_0_move_energy=collector.sensor_states["gate_0_move_energy"],
+            light=collector.sensor_states["light"],
+        )
 
         # Verify the truncated frame warning was logged
         assert truncated_warning_seen.done(), (
@@ -490,13 +483,12 @@ async def test_uart_mock_ld2412_footer_in_payload(
 
         motion_written = events.next_line_with(expected_motion)
         client.number_command(numbers["gate_0_move_threshold"].key, 77.0)
-        try:
-            await asyncio.wait_for(motion_written, timeout=5.0)
-        except TimeoutError:
-            pytest.fail(
-                "Timeout waiting for the motion threshold command.\n"
-                f"  expected motion payload: {expected_motion}"
-            )
+        await wait_or_fail(
+            motion_written,
+            5.0,
+            "Timeout waiting for the motion threshold command",
+            expected_payload=expected_motion,
+        )
 
         assert events.count(STILL_COMMAND) == 0, (
             "still threshold command was sent for a motion write"
@@ -504,13 +496,12 @@ async def test_uart_mock_ld2412_footer_in_payload(
 
         still_written = events.next_line_with(expected_still)
         client.number_command(numbers["gate_5_still_threshold"].key, 77.0)
-        try:
-            await asyncio.wait_for(still_written, timeout=5.0)
-        except TimeoutError:
-            pytest.fail(
-                "Timeout waiting for the still threshold command.\n"
-                f"  expected still payload: {expected_still}"
-            )
+        await wait_or_fail(
+            still_written,
+            5.0,
+            "Timeout waiting for the still threshold command",
+            expected_payload=expected_still,
+        )
 
         assert events.count(MOTION_COMMAND) == 1
         assert events.count(STILL_COMMAND) == 1
@@ -519,24 +510,20 @@ async def test_uart_mock_ld2412_footer_in_payload(
         # The reset discards the module's thresholds, so the next write is held back and re-queried
         reset_sent = events.next_line_with(FACTORY_RESET)
         client.button_command(require_entity(entities, "factory_reset", ButtonInfo).key)
-        try:
-            await asyncio.wait_for(reset_sent, timeout=5.0)
-        except TimeoutError:
-            pytest.fail("Timeout waiting for the factory reset command")
+        await wait_or_fail(
+            reset_sent, 5.0, "Timeout waiting for the factory reset command"
+        )
 
         held_back = events.next_line_with(MOTION_HELD_BACK)
         written_after_reset = events.next_line_with(expected_motion_after_reset)
         client.number_command(numbers["gate_0_move_threshold"].key, 50.0)
-        try:
-            await asyncio.wait_for(
-                asyncio.gather(held_back, written_after_reset), timeout=5.0
-            )
-        except TimeoutError:
-            pytest.fail(
-                "Timeout waiting for the write after the factory reset.\n"
-                f"  expected motion payload: {expected_motion_after_reset}\n"
-                f"  logged: {events.lines}"
-            )
+        await wait_or_fail(
+            asyncio.gather(held_back, written_after_reset),
+            5.0,
+            "Timeout waiting for the write after the factory reset",
+            expected_payload=expected_motion_after_reset,
+            logged=events.lines,
+        )
         reset_at = events.index_after(-1, FACTORY_RESET)
         warning_at = events.index_after(reset_at, MOTION_HELD_BACK)
         query_at = events.index_after(warning_at, MOTION_QUERY)
@@ -658,12 +645,12 @@ async def test_uart_mock_ld2412_requery_after_giveup(
             for key in (gate_0_key, gate_5_key)
         ]
         client.number_command(gate_5_key, 20.0)
-        try:
-            await asyncio.wait_for(
-                asyncio.gather(held_back, config_mode_left), timeout=5.0
-            )
-        except TimeoutError:
-            pytest.fail(f"Timeout waiting for the write to be handled: {events.lines}")
+        await wait_or_fail(
+            asyncio.gather(held_back, config_mode_left),
+            5.0,
+            "Timeout waiting for the write to be handled",
+            logged=events.lines,
+        )
         # The unanswered query times out after 2 s
         try:
             await asyncio.wait_for(given_up, timeout=3.0)

tests/integration/test_uart_mock_ld2420.py

@@ -24,8 +24,9 @@ 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
+from .wait_utils import wait_or_fail
 
 
 @pytest.mark.asyncio
@@ -86,8 +87,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
@@ -106,13 +106,12 @@ async def test_uart_mock_ld2420(
 
         # Wait for the recovery frame (Phase 5) to be parsed
         # This proves the component survived garbage + truncated + overflow
-        try:
-            await asyncio.wait_for(recovery_received, timeout=5.0)
-        except TimeoutError:
-            pytest.fail(
-                f"Timeout waiting for recovery frame. Received:\n"
-                f"  sensor_states: {collector.sensor_states}"
-            )
+        await wait_or_fail(
+            recovery_received,
+            5.0,
+            "Timeout waiting for recovery frame",
+            sensor_states=collector.sensor_states,
+        )
 
         # Verify overflow warning was logged
         assert overflow_seen.done(), (
@@ -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
@@ -217,13 +215,12 @@ async def test_uart_mock_ld2420_simple(
         assert collector.binary_states["has_target"][0] is True
 
         # Wait for Phase 4 recovery (distance=50) after overflow
-        try:
-            await asyncio.wait_for(recovery_received, timeout=5.0)
-        except TimeoutError:
-            pytest.fail(
-                f"Timeout waiting for recovery frame. Received:\n"
-                f"  moving_distance: {collector.sensor_states['moving_distance']}"
-            )
+        await wait_or_fail(
+            recovery_received,
+            5.0,
+            "Timeout waiting for recovery frame",
+            moving_distance=collector.sensor_states["moving_distance"],
+        )
 
         # Verify overflow warning was logged
         assert overflow_seen.done(), (
@@ -234,16 +231,15 @@ async def test_uart_mock_ld2420_simple(
         # This assertion FAILS before PR #14458 because the 16-digit frame
         # in Phase 5 causes an infinite loop in handle_simple_mode_ pre-fix.
         # The binary hangs, Phase 6 never fires, and this wait times out.
-        try:
-            await asyncio.wait_for(post_bug_received, timeout=8.0)
-        except TimeoutError:
-            pytest.fail(
-                f"Timeout waiting for post-bug recovery (distance=25). "
-                f"This likely means Phase 5 (16-digit frame) caused an infinite "
-                f"loop in handle_simple_mode_, indicating PR #14458 bug #1 fix "
-                f"is missing.\n"
-                f"  moving_distance values: {collector.sensor_states['moving_distance']}"
-            )
+        await wait_or_fail(
+            post_bug_received,
+            8.0,
+            "Timeout waiting for post-bug recovery (distance=25). "
+            "This likely means Phase 5 (16-digit frame) caused an infinite "
+            "loop in handle_simple_mode_, indicating PR #14458 bug #1 fix "
+            "is missing.",
+            moving_distance=collector.sensor_states["moving_distance"],
+        )
 
         # Verify post-bug value
         post_bug_values = [

tests/integration/test_uart_mock_ld2420_config_exit.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import asyncio
 import re
 
 from aioesphomeapi import ButtonInfo
@@ -10,6 +9,7 @@ import pytest
 
 from .state_utils import InitialStateHelper, SensorStateCollector, require_entity
 from .types import APIClientConnectedFactory, RunCompiledFunction
+from .wait_utils import wait_or_fail
 
 EXIT_SENT = re.compile(r"Config exit sent (\d+)")
 
@@ -62,14 +62,13 @@ async def test_uart_mock_ld2420_config_exit(
         resumed = collector.add_waiter(
             lambda: len(collector.sensor_states["moving_distance"]) >= 2
         )
-        try:
-            await asyncio.wait_for(resumed, timeout=5.0)
-        except TimeoutError:
-            pytest.fail(
-                "Energy reports did not resume after the last config exit.\n"
-                f"  exit sends: {exit_sends}\n"
-                f"  sensor_states: {collector.sensor_states}"
-            )
+        await wait_or_fail(
+            resumed,
+            5.0,
+            "Energy reports did not resume after the last config exit",
+            exit_sends=exit_sends,
+            sensor_states=collector.sensor_states,
+        )
 
         # The sensor skips repeated values, so each distinct distance is one report
         assert collector.sensor_states["moving_distance"] == pytest.approx(

tests/integration/test_uart_mock_ld2450.py

@@ -22,8 +22,9 @@ 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
+from .wait_utils import wait_or_fail
 
 
 @pytest.mark.asyncio
@@ -108,8 +109,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
@@ -159,13 +159,12 @@ async def test_uart_mock_ld2450(
 
         # Wait for the recovery frame (Phase 5) to be parsed
         # This proves the component survived garbage + truncated + overflow
-        try:
-            await asyncio.wait_for(recovery_received, timeout=5.0)
-        except TimeoutError:
-            pytest.fail(
-                f"Timeout waiting for recovery frame. Received:\n"
-                f"  sensor_states: {collector.sensor_states}"
-            )
+        await wait_or_fail(
+            recovery_received,
+            5.0,
+            "Timeout waiting for recovery frame",
+            sensor_states=collector.sensor_states,
+        )
 
         # Verify overflow warning was logged
         assert overflow_seen.done(), (

tests/integration/test_uart_mock_ld2460.py

@@ -19,8 +19,9 @@ 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
+from .wait_utils import wait_or_fail
 
 
 @pytest.mark.asyncio
@@ -89,8 +90,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)
@@ -101,13 +101,12 @@ async def test_uart_mock_ld2460(
                 in collector.sensor_states["target_1_angle"]
             )
         )
-        try:
-            await asyncio.wait_for(phase1, timeout=5.0)
-        except TimeoutError:
-            pytest.fail(
-                f"Timeout waiting for the phase 1 frame. Received:\n"
-                f"  sensor_states: {collector.sensor_states}"
-            )
+        await wait_or_fail(
+            phase1,
+            5.0,
+            "Timeout waiting for the phase 1 frame",
+            sensor_states=collector.sensor_states,
+        )
         assert collector.sensor_states["target_1_x"][0] == pytest.approx(1.5, abs=0.01)
         assert collector.sensor_states["target_1_y"][0] == pytest.approx(2.3, abs=0.01)
         assert collector.sensor_states["target_1_distance"][0] == pytest.approx(
@@ -119,13 +118,12 @@ async def test_uart_mock_ld2460(
         assert collector.sensor_states["target_count"][0] == pytest.approx(1.0)
 
         # Phase 5: parsing survives garbage, a truncated frame and an overflow
-        try:
-            await asyncio.wait_for(recovery_received, timeout=5.0)
-        except TimeoutError:
-            pytest.fail(
-                f"Timeout waiting for the recovery frame. Received:\n"
-                f"  sensor_states: {collector.sensor_states}"
-            )
+        await wait_or_fail(
+            recovery_received,
+            5.0,
+            "Timeout waiting for the recovery frame",
+            sensor_states=collector.sensor_states,
+        )
         assert overflow_seen.done(), "Expected 'Max frame length exceeded' in the log"
         assert (
             pytest.approx(3.606, abs=0.01)
@@ -137,13 +135,12 @@ async def test_uart_mock_ld2460(
         assert pytest.approx(26.6, abs=0.1) in collector.sensor_states["target_2_angle"]
 
         # Phase 6: the top-mount receipt changes the angle for the same coordinates
-        try:
-            await asyncio.wait_for(top_mount_received, timeout=5.0)
-        except TimeoutError:
-            pytest.fail(
-                f"Timeout waiting for the top-mount angle. Received:\n"
-                f"  sensor_states: {collector.sensor_states}"
-            )
+        await wait_or_fail(
+            top_mount_received,
+            5.0,
+            "Timeout waiting for the top-mount angle",
+            sensor_states=collector.sensor_states,
+        )
 
         # The hub sent the installation-mode query from the datasheet (table 11)
         tx_data = " ".join(tx_log_lines)

tests/integration/test_uart_mock_modbus.py

@@ -23,8 +23,9 @@ 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
+from .wait_utils import wait_or_fail
 
 
 def _swap16(value: int) -> int:
@@ -832,10 +833,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 +889,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 +927,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 +1047,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 +1062,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 +1102,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 +1138,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 +1195,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)
@@ -1289,8 +1281,5 @@ async def test_uart_mock_modbus_ranges(
         await tracker.setup_and_start_scenario(client)
         await tracker.await_all(futures)
         # text_block is not tracker-registered (non-numeric), so time out explicitly.
-        try:
-            await asyncio.wait_for(text_future, timeout=5.0)
-        except TimeoutError:
-            pytest.fail("text_block never published 'ABCDEF'")
+        await wait_or_fail(text_future, 5.0, "text_block never published 'ABCDEF'")
         _assert_no_modbus_errors(error_log_lines, warning_log_lines)

tests/integration/test_uart_mock_pylontech.py

@@ -16,8 +16,9 @@ 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
+from .wait_utils import wait_or_fail
 
 SENSORS = [
     "bat1_voltage",
@@ -102,18 +103,16 @@ 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:
-            await asyncio.wait_for(recovery_received, timeout=5.0)
-        except TimeoutError:
-            pytest.fail(
-                f"Timeout waiting for the recovery reply. Received:\n"
-                f"  sensor_states: {collector.sensor_states}\n"
-                f"  text_states: {collector.text_sensor_states}"
-            )
+        await wait_or_fail(
+            recovery_received,
+            5.0,
+            "Timeout waiting for the recovery reply",
+            sensor_states=collector.sensor_states,
+            text_states=collector.text_sensor_states,
+        )
         # Let any stray publish from the bad phases arrive before the exact checks
         await asyncio.sleep(0.2)
 

tests/integration/test_uart_mock_ufm01.py

@@ -40,6 +40,7 @@ import pytest
 
 from .state_utils import InitialStateHelper, StateWaiter
 from .types import APIClientConnectedFactory, RunCompiledFunction
+from .wait_utils import wait_or_fail
 
 # Commands as the uart mock logs them: FE FE 11, command, argument, checksum, 16
 RESET_DEVICE = "FE:FE:11:5D:CB:28:16"
@@ -199,10 +200,7 @@ async def _gather(meter: _Meter, waits: list[Awaitable[EntityState]]) -> None:
 
 
 async def _wait_for(log: _Log, future: asyncio.Future[float], timeout: float) -> float:
-    try:
-        return await asyncio.wait_for(future, timeout)
-    except TimeoutError:
-        pytest.fail(f"Timed out; log: {log.events}")
+    return await wait_or_fail(future, timeout, "Timed out", log=log.events)
 
 
 @pytest.mark.shared_yaml("uart_mock_ufm01")

tests/integration/test_wait_until_on_boot.py

@@ -19,6 +19,7 @@ import pytest
 
 from .state_utils import require_service
 from .types import APIClientConnectedFactory, RunCompiledFunction
+from .wait_utils import wait_or_fail
 
 
 @pytest.mark.asyncio
@@ -76,14 +77,13 @@ async def test_wait_until_on_boot(
 
         # If the fix works, wait_until's loop() will check the condition and proceed
         # If the bug exists, wait_until is stuck with disabled loop and will timeout
-        try:
-            await asyncio.wait_for(on_boot_complete_future, timeout=2.0)
-            assert on_boot_completed, (
-                "on_boot wait_until did not complete after flag was set"
-            )
-        except TimeoutError:
-            pytest.fail(
-                "wait_until in on_boot did not complete within 2s after condition became true. "
-                "This indicates the bug where WaitUntilAction::setup() disables the loop "
-                "after play_complex() has already enabled it."
-            )
+        await wait_or_fail(
+            on_boot_complete_future,
+            2.0,
+            "wait_until in on_boot did not complete within 2s after condition became true. "
+            "This indicates the bug where WaitUntilAction::setup() disables the loop "
+            "after play_complex() has already enabled it.",
+        )
+        assert on_boot_completed, (
+            "on_boot wait_until did not complete after flag was set"
+        )

tests/integration/test_wake_loop_forces_phase_b.py

@@ -26,6 +26,7 @@ import pytest
 
 from .state_utils import assert_device_name
 from .types import APIClientConnectedFactory, RunCompiledFunction
+from .wait_utils import wait_or_fail
 
 
 @pytest.mark.asyncio
@@ -51,10 +52,9 @@ async def test_wake_loop_forces_phase_b(
     ):
         await assert_device_name(client, "wake-loop-phase-b")
 
-        try:
-            delta, elapsed = await asyncio.wait_for(result, timeout=15.0)
-        except TimeoutError:
-            pytest.fail("WAKE_RESULT marker never appeared")
+        delta, elapsed = await wait_or_fail(
+            result, 15.0, "WAKE_RESULT marker never appeared"
+        )
 
         # Without the fix, delta would be 0 — loop_interval_=2000ms held
         # Phase B off for the full 500ms observation window. With the fix