Let wait_for_initial_states report the missing entities itself

# Conflicts:
#	tests/integration/README.md

tests/integration/README.md

@@ -58,8 +58,30 @@ The `InitialStateHelper` class solves a common problem in integration tests: whe
 - Provides debug logging to track state reception (use `--log-cli-level=DEBUG`)
 - Safe for concurrent use with multiple entity types
 
-**Future work:**
-Consider converting existing integration tests to use `InitialStateHelper` for more reliable state tracking and to eliminate race conditions related to initial state broadcasts.
+#### Subscribing after the initial states (`state_utils.py`)
+
+`subscribe_after_initial_states()` does the `InitialStateHelper` setup in one call: it subscribes, waits
+for every entity's connect time state, and fails the test with the entities still missing on timeout.
+It returns a `StateSubscription` with `initial_states` (by key) and a `waiter`, a `StateWaiter` that sees
+every later state. `waiter.next_state(key=..., state_type=...)` arms a wait for the next matching state;
+call it before sending the command. `press_and_wait()` presses a button and returns the next state of a
+type, from one entity when `key=` is given:
+
+```python
+from .state_utils import press_and_wait, subscribe_after_initial_states
+
+entities, _ = await client.list_entities_services()
+sub = await subscribe_after_initial_states(client, entities)
+
+state = await press_and_wait(client, entities, sub.waiter, "set_position", CoverState)
+assert state.position == pytest.approx(0.5, abs=0.01)
+
+next_light = sub.waiter.next_state(key=light.key, state_type=LightState)
+client.light_command(key=light.key, state=True)
+assert (await next_light).state is True
+```
+
+Pass `on_state=` as well when the test also needs to collect every state.
 
 #### Entity Lookup Helpers (`state_utils.py`)
 
@@ -203,47 +225,24 @@ async def test_my_sensor(
 
 ##### State Subscription Pattern
 
-**Recommended: Using InitialStateHelper**
+**Recommended: Using subscribe_after_initial_states**
 
-When an API client connects, ESPHome automatically sends the current state of all entities. The `InitialStateHelper` (from `state_utils.py`) handles this by swallowing these initial states and only forwarding subsequent state changes to your test callback:
+When an API client connects, ESPHome automatically sends the current state of all entities.
+`subscribe_after_initial_states()` (from `state_utils.py`) swallows these initial states and only forwards
+subsequent state changes, to your callback and to the returned `waiter`:
 
 ```python
-from .state_utils import InitialStateHelper
+from .state_utils import subscribe_after_initial_states
 
-# Track state changes with futures
-loop = asyncio.get_running_loop()
-states: dict[int, EntityState] = {}
-state_future: asyncio.Future[EntityState] = loop.create_future()
-
-
-def on_state(state: EntityState) -> None:
-    """This callback only receives NEW state changes, not initial states."""
-    states[state.key] = state
-    # Check for specific condition using isinstance
-    if isinstance(state, SensorState) and state.state == expected_value:
-        if not state_future.done():
-            state_future.set_result(state)
-
-
-# Get entities and set up state synchronization
 entities, services = await client.list_entities_services()
-initial_state_helper = InitialStateHelper(entities)
+sub = await subscribe_after_initial_states(client, 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; 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
+# Arm the wait before the action that produces the state
+expected = sub.waiter.expect(
+    lambda s: isinstance(s, SensorState) and s.state == expected_value
+)
 # ... trigger state changes ...
-
-# Wait for expected state
-try:
-    result = await asyncio.wait_for(state_future, timeout=5.0)
-except asyncio.TimeoutError:
-    pytest.fail(f"Expected state not received. Got: {list(states.values())}")
+result = await expected
 ```
 
 **Legacy: Manual State Tracking**

tests/integration/state_utils.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import asyncio
 from collections.abc import Awaitable, Callable
+from dataclasses import dataclass
 import logging
 from typing import TypeVar
 
@@ -100,6 +101,22 @@ class StateWaiter:
         self._waiters.append(entry)
         return self._wait(entry, timeout, label)
 
+    def next_state[T: EntityState](
+        self,
+        key: int | None = None,
+        state_type: type[T] = EntityState,
+        timeout: float = 5.0,
+    ) -> Awaitable[T]:
+        """Arm a wait for the next ``state_type`` state, for entity ``key`` if given.
+
+        Like :meth:`expect`, call it before sending the command and await the result afterwards.
+        """
+        return self.expect(
+            lambda s: isinstance(s, state_type) and (key is None or s.key == key),
+            timeout,
+            state_type.__name__ if key is None else f"{state_type.__name__} key={key}",
+        )
+
     async def _wait(
         self,
         entry: tuple[Callable[[EntityState], bool], asyncio.Future[EntityState]],
@@ -382,6 +399,13 @@ class InitialStateHelper:
 
         return wrapper
 
+    @property
+    def pending_object_ids(self) -> list[str]:
+        """Object ids of the entities whose initial state has not arrived yet."""
+        return sorted(
+            self._entities_by_id[k].object_id for k in self._wait_initial_states
+        )
+
     async def wait_for_initial_states(self, timeout: float = 5.0) -> None:
         """Wait for all initial states to be received, failing the test on timeout.
 
@@ -402,6 +426,58 @@ class InitialStateHelper:
         )
 
 
+@dataclass(frozen=True, slots=True)
+class StateSubscription:
+    """A state subscription that is past the connect time states."""
+
+    initial_states: dict[int, EntityState]
+    waiter: StateWaiter
+
+
+async def subscribe_after_initial_states(
+    client: APIClient,
+    entities: list[EntityInfo],
+    on_state: Callable[[EntityState], None] | None = None,
+    timeout: float = 5.0,
+) -> StateSubscription:
+    """Subscribe to states and wait until every entity has sent its connect time state.
+
+    Later states go to ``on_state`` (if given) and to the returned ``waiter``. The connect time
+    states are kept in ``initial_states`` by key. Fails the test listing the entities still missing.
+    """
+    helper = InitialStateHelper(entities)
+    waiter = StateWaiter()
+
+    def forward(state: EntityState) -> None:
+        waiter.on_state(state)
+        if on_state is not None:
+            on_state(state)
+
+    client.subscribe_states(helper.on_state_wrapper(forward))
+    await helper.wait_for_initial_states(timeout)
+    return StateSubscription(helper.initial_states, waiter)
+
+
+async def press_and_wait[T: EntityState](
+    client: APIClient,
+    entities: list[EntityInfo],
+    waiter: StateWaiter,
+    button: str,
+    state_type: type[T],
+    *,
+    key: int | None = None,
+    timeout: float = 5.0,
+) -> T:
+    """Press the button whose object_id contains ``button`` and return the next ``state_type`` state.
+
+    Pass the entity's ``key`` when the fixture has more than one entity of that type.
+    """
+    btn = require_entity(entities, button, ButtonInfo)
+    next_state = waiter.next_state(key=key, state_type=state_type, timeout=timeout)
+    client.button_command(btn.key)
+    return await next_state
+
+
 class SensorStateCollector:
     """Collects sensor, binary sensor, and text sensor state updates with wait helpers.
 

tests/unit_tests/test_integration_state_subscribe.py

@@ -0,0 +1,108 @@
+"""Tests for the state subscription helpers in tests/integration/state_utils.py."""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Callable
+
+from aioesphomeapi import ButtonInfo, EntityState, SensorInfo, SensorState, SwitchState
+import pytest
+
+from tests.integration.state_utils import (
+    StateSubscription,
+    StateWaiter,
+    press_and_wait,
+    subscribe_after_initial_states,
+)
+
+ENTITIES = [
+    SensorInfo(object_id="temp", key=1),
+    SensorInfo(object_id="humidity", key=2),
+    ButtonInfo(object_id="set_temp", key=3),
+]
+
+
+class FakeClient:
+    """Records the state callback and button presses."""
+
+    def __init__(self) -> None:
+        self.callback: Callable[[EntityState], None] | None = None
+        self.pressed: list[int] = []
+
+    def subscribe_states(self, callback: Callable[[EntityState], None]) -> None:
+        self.callback = callback
+
+    def button_command(self, key: int) -> None:
+        self.pressed.append(key)
+        self.callback(SensorState(key=1, state=42.0))
+
+
+async def _subscribe(client: FakeClient, **kwargs: object) -> StateSubscription:
+    task = asyncio.create_task(
+        subscribe_after_initial_states(client, ENTITIES, **kwargs)
+    )
+    await asyncio.sleep(0)
+    client.callback(SensorState(key=1, state=1.0))
+    client.callback(SensorState(key=2, state=2.0))
+    return await task
+
+
+@pytest.mark.asyncio
+async def test_subscribe_swallows_initial_states_and_forwards_later_ones() -> None:
+    client = FakeClient()
+    received: list[EntityState] = []
+    sub = await _subscribe(client, on_state=received.append)
+    assert sub.initial_states[1].state == 1.0
+    assert sub.initial_states[2].state == 2.0
+    assert received == []
+
+    next_state = sub.waiter.next_state(key=2, state_type=SensorState)
+    client.callback(SensorState(key=1, state=3.0))
+    client.callback(SensorState(key=2, state=4.0))
+    assert (await next_state).state == 4.0
+    assert [s.state for s in received] == [3.0, 4.0]
+
+
+@pytest.mark.asyncio
+async def test_subscribe_timeout_lists_missing_entities() -> None:
+    client = FakeClient()
+    task = asyncio.create_task(
+        subscribe_after_initial_states(client, ENTITIES, timeout=0.01)
+    )
+    await asyncio.sleep(0)
+    client.callback(SensorState(key=1, state=1.0))
+    with pytest.raises(pytest.fail.Exception, match=r"still_waiting[\s\S]*humidity"):
+        await task
+
+
+@pytest.mark.asyncio
+async def test_next_state_filters_by_type_and_times_out() -> None:
+    waiter = StateWaiter()
+    next_switch = waiter.next_state(state_type=SwitchState)
+    waiter.on_state(SensorState(key=1, state=1.0))
+    waiter.on_state(SwitchState(key=5, state=True))
+    assert (await next_switch).key == 5
+    with pytest.raises(TimeoutError, match="SwitchState key=9"):
+        await waiter.next_state(key=9, state_type=SwitchState, timeout=0.01)
+
+
+@pytest.mark.asyncio
+async def test_press_and_wait_presses_button_and_returns_state() -> None:
+    client = FakeClient()
+    sub = await _subscribe(client)
+    state = await press_and_wait(client, ENTITIES, sub.waiter, "set_temp", SensorState)
+    assert client.pressed == [3]
+    assert state.state == 42.0
+
+
+@pytest.mark.asyncio
+async def test_press_and_wait_waits_for_the_given_key() -> None:
+    client = FakeClient()
+    sub = await _subscribe(client)
+    pending = asyncio.ensure_future(
+        press_and_wait(client, ENTITIES, sub.waiter, "set_temp", SensorState, key=2)
+    )
+    await asyncio.sleep(0)
+    assert not pending.done()  # the press published key 1
+    client.callback(SensorState(key=2, state=55.0))
+    assert (await pending).state == 55.0