[tests] Fail the test when the integration state helpers time out (#21244)

tests/integration/README.md

@@ -45,7 +45,7 @@ The `InitialStateHelper` class solves a common problem in integration tests: whe
 - Tracks all entities (except stateless ones like buttons)
 - Swallows the first state broadcast for each entity
 - Forwards all subsequent state changes to your test callback
-- Provides `wait_for_initial_states()` to synchronize before test actions
+- Provides `wait_for_initial_states()` to synchronize before test actions; it fails the test on timeout, naming the entities still without a state
 
 **When to use it:**
 - Any test that triggers entity state changes and needs to verify them
@@ -232,11 +232,9 @@ initial_state_helper = InitialStateHelper(entities)
 # Subscribe with the wrapper that filters initial states
 client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
 
-# Wait for all initial states to be broadcast
-try:
-    await initial_state_helper.wait_for_initial_states()
-except TimeoutError:
-    pytest.fail("Timeout waiting for initial states")
+# Wait for all initial states to be broadcast; on timeout this fails the test
+# and lists the entities that never sent one
+await initial_state_helper.wait_for_initial_states()
 
 # Now perform your test actions - on_state will only receive new changes
 # ... trigger state changes ...

tests/integration/state_utils.py

@@ -19,6 +19,8 @@ from aioesphomeapi import (
     UserService,
 )
 
+from .wait_utils import wait_or_fail
+
 _LOGGER = logging.getLogger(__name__)
 
 S = TypeVar("S", bound=EntityState)
@@ -28,6 +30,7 @@ async def wait_for_state(
     client: APIClient,
     predicate: Callable[[EntityState], bool],
     timeout: float = 5.0,
+    msg: str = "Timed out waiting for a matching state",
 ) -> EntityState:
     """Subscribe to states and wait for one matching ``predicate``.
 
@@ -42,21 +45,24 @@ async def wait_for_state(
         predicate: Callable invoked for every received state; the first
             state for which it returns ``True`` is returned.
         timeout: Maximum time to wait in seconds.
+        msg: Failure message if no matching state arrives in time; the
+            states received are listed after it.
 
     Returns:
         The first state matching ``predicate``.
-
-    Raises:
-        asyncio.TimeoutError: If no matching state arrives within ``timeout``.
     """
     future: asyncio.Future[EntityState] = asyncio.get_running_loop().create_future()
+    received: list[EntityState] = []
 
     def on_state(state: EntityState) -> None:
-        if not future.done() and predicate(state):
+        if future.done():
+            return
+        received.append(state)
+        if predicate(state):
             future.set_result(state)
 
     client.subscribe_states(on_state)
-    return await asyncio.wait_for(future, timeout=timeout)
+    return await wait_or_fail(future, timeout, msg, received=received)
 
 
 class StateWaiter:
@@ -377,15 +383,23 @@ class InitialStateHelper:
         return wrapper
 
     async def wait_for_initial_states(self, timeout: float = 5.0) -> None:
-        """Wait for all initial states to be received.
+        """Wait for all initial states to be received, failing the test on timeout.
 
         Args:
             timeout: Maximum time to wait in seconds
-
-        Raises:
-            asyncio.TimeoutError: If initial states aren't received within timeout
         """
-        await asyncio.wait_for(self._initial_states_received, timeout=timeout)
+        await wait_or_fail(
+            self._initial_states_received,
+            timeout,
+            "Timed out waiting for initial states",
+            still_waiting=lambda: sorted(
+                f"{type(entity).__name__} {entity.object_id}"
+                for entity in (
+                    self._entities_by_id[entity_id]
+                    for entity_id in self._wait_initial_states
+                )
+            ),
+        )
 
 
 class SensorStateCollector:
@@ -471,12 +485,29 @@ class SensorStateCollector:
         )
 
     async def wait_for_all(self, timeout: float = 3.0) -> None:
-        """Wait until all sensors and binary sensors have at least one value."""
+        """Wait until every tracked sensor has a value, failing the test on timeout."""
         if self._all_have_values():
             return
         future: asyncio.Future[bool] = asyncio.get_running_loop().create_future()
         self._waiters.append((self._all_have_values, future))
-        await asyncio.wait_for(future, timeout=timeout)
+        tracked = (self.sensor_states, self.binary_states, self.text_sensor_states)
+        await wait_or_fail(
+            future,
+            timeout,
+            "Timed out waiting for every sensor to report",
+            still_waiting=lambda: [
+                name
+                for states in tracked
+                for name, values in states.items()
+                if not values
+            ],
+            received=lambda: {
+                name: values
+                for states in tracked
+                for name, values in states.items()
+                if values
+            },
+        )
 
     def add_waiter(self, condition: Callable[[], bool]) -> asyncio.Future[bool]:
         """Add a custom waiter that resolves when condition returns True.

tests/integration/wait_utils.py

@@ -16,6 +16,8 @@ def _format_context(context: dict[str, object]) -> str:
     """Render context values for a failure message, keeping only the tail of collections."""
     parts: list[str] = []
     for name, value in context.items():
+        if callable(value):
+            value = value()
         if isinstance(value, Mapping):
             items = [f"{key}: {item!r}" for key, item in value.items()]
         elif isinstance(value, (list, tuple, deque)):
@@ -39,8 +41,8 @@ async def wait_or_fail[T](
 
     Behaves like ``asyncio.wait_for``, so a timed out future is cancelled; a TimeoutError raised
     by ``aw`` itself propagates unchanged, and ``timeout`` must be positive. ``context`` values
-    are formatted only on failure, so pass live collections; lists, tuples, deques and mappings
-    show their last 20 items.
+    are formatted only on failure, so pass live collections, or a zero-argument callable to
+    compute a value then; lists, tuples, deques and mappings show their last 20 items.
     """
     if timeout <= 0:
         # asyncio.wait_for treats this as already timed out; a test should never ask for it

tests/unit_tests/test_integration_state_utils.py

@@ -2,25 +2,31 @@
 
 from __future__ import annotations
 
-from unittest.mock import AsyncMock
+import asyncio
+from unittest.mock import AsyncMock, MagicMock
 
 from aioesphomeapi import (
     BinarySensorInfo,
+    BinarySensorState,
     ButtonInfo,
     DeviceInfo,
     EntityInfo,
     LightInfo,
     SensorInfo,
+    SensorState,
     UserService,
 )
 import pytest
 
 from tests.integration.state_utils import (
+    InitialStateHelper,
+    SensorStateCollector,
     assert_device_name,
     find_entity,
     require_entity,
     require_only_entity,
     require_service,
+    wait_for_state,
 )
 
 ENTITIES: list[EntityInfo] = [
@@ -75,3 +81,48 @@ async def test_assert_device_name() -> None:
     client.device_info.return_value = None
     with pytest.raises(AssertionError, match="returned None"):
         await assert_device_name(client, "host-test")
+
+
+@pytest.mark.asyncio
+async def test_wait_for_initial_states_fails_naming_missing_entities() -> None:
+    helper = InitialStateHelper(ENTITIES)
+    wrapper = helper.on_state_wrapper(lambda state: None)
+    wrapper(SensorState(key=1, state=1.0))
+    wrapper(BinarySensorState(key=3, state=True))
+    with pytest.raises(pytest.fail.Exception) as err:
+        await helper.wait_for_initial_states(timeout=0.01)
+    assert str(err.value) == (
+        "Timed out waiting for initial states (after 0.01s)\n"
+        "still_waiting (last 1 of 1):\n  SensorInfo temperature"
+    )
+
+
+@pytest.mark.asyncio
+async def test_wait_for_all_fails_naming_missing_sensors() -> None:
+    collector = SensorStateCollector(
+        ["temperature_raw", "temperature"], ["Motion"], entities=ENTITIES
+    )
+    collector.on_state(SensorState(key=1, state=1.5))
+    with pytest.raises(pytest.fail.Exception) as err:
+        await collector.wait_for_all(timeout=0.01)
+    message = str(err.value)
+    assert message.startswith("Timed out waiting for every sensor to report")
+    assert "still_waiting (last 2 of 2):\n  temperature\n  Motion" in message
+    assert message.endswith("received (last 1 of 1):\n  temperature_raw: [1.5]")
+
+
+@pytest.mark.asyncio
+async def test_wait_for_state_fails_listing_received_states() -> None:
+    client = MagicMock()
+    waiting = asyncio.ensure_future(
+        wait_for_state(client, lambda state: state.state == 2.0, 0.01, "no 2.0")
+    )
+    await asyncio.sleep(0)
+    on_state = client.subscribe_states.call_args.args[0]
+    on_state(SensorState(key=1, state=1.0))
+    with pytest.raises(pytest.fail.Exception) as err:
+        await waiting
+    message = str(err.value)
+    assert message.startswith("no 2.0 (after 0.01s)\nreceived (last 1 of 1):")
+    assert "state=1.0" in message
+    on_state(SensorState(key=1, state=2.0))

tests/unit_tests/test_integration_wait_utils.py

@@ -71,6 +71,18 @@ async def test_wait_or_fail_trims_deques_and_mappings() -> None:
     assert "pairs (last 20 of 25):\n  5\n" in message
 
 
+@pytest.mark.asyncio
+async def test_wait_or_fail_calls_callable_context_on_failure() -> None:
+    pending = ["a", "b"]
+    waiting = asyncio.ensure_future(
+        wait_or_fail(asyncio.Event(), 0.01, "stuck", missing=lambda: sorted(pending))
+    )
+    pending.remove("a")
+    with pytest.raises(pytest.fail.Exception) as err:
+        await waiting
+    assert str(err.value).endswith("missing (last 1 of 1):\n  b")
+
+
 @pytest.mark.asyncio
 async def test_wait_or_fail_reraises_a_timeout_from_the_awaitable() -> None:
     async def inner() -> None: