[tests] Fix flaky host_mode_climate_basic_state integration test (#16192)

tests/integration/fixtures/host_mode_climate_basic_state.yaml

@@ -1,5 +1,5 @@
 esphome:
-  name: host-climate-test
+  name: host-climate-basic-state
 host:
 api:
 logger:
@@ -10,6 +10,7 @@ climate:
     name: Dual-mode Thermostat
     sensor: host_thermostat_temperature_sensor
     humidity_sensor: host_thermostat_humidity_sensor
+    on_boot_restore_from: default_preset
     humidity_hysteresis: 1.0
     min_cooling_off_time: 20s
     min_cooling_run_time: 20s

tests/integration/state_utils.py

@@ -8,6 +8,7 @@ import logging
 from typing import TypeVar
 
 from aioesphomeapi import (
+    APIClient,
     BinarySensorState,
     ButtonInfo,
     EntityInfo,
@@ -19,6 +20,42 @@ from aioesphomeapi import (
 _LOGGER = logging.getLogger(__name__)
 
 T = TypeVar("T", bound=EntityInfo)
+S = TypeVar("S", bound=EntityState)
+
+
+async def wait_for_state(
+    client: APIClient,
+    predicate: Callable[[EntityState], bool],
+    timeout: float = 5.0,
+) -> EntityState:
+    """Subscribe to states and wait for one matching ``predicate``.
+
+    Resolves with the first :class:`EntityState` for which ``predicate``
+    returns ``True``. Useful when a component publishes multiple states
+    during setup (e.g. before sensor readings arrive) and the test needs
+    to wait for the state to converge to expected values rather than
+    capturing whichever state happens to arrive first.
+
+    Args:
+        client: Connected API client.
+        predicate: Callable invoked for every received state; the first
+            state for which it returns ``True`` is returned.
+        timeout: Maximum time to wait in seconds.
+
+    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()
+
+    def on_state(state: EntityState) -> None:
+        if not future.done() and predicate(state):
+            future.set_result(state)
+
+    client.subscribe_states(on_state)
+    return await asyncio.wait_for(future, timeout=timeout)
 
 
 def find_entity(

tests/integration/test_host_mode_climate_basic_state.py

@@ -2,11 +2,17 @@
 
 from __future__ import annotations
 
-import aioesphomeapi
-from aioesphomeapi import ClimateAction, ClimateInfo, ClimateMode, ClimatePreset
+from aioesphomeapi import (
+    ClimateAction,
+    ClimateInfo,
+    ClimateMode,
+    ClimatePreset,
+    ClimateState,
+    EntityState,
+)
 import pytest
 
-from .state_utils import InitialStateHelper
+from .state_utils import wait_for_state
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -18,32 +24,30 @@ async def test_host_mode_climate_basic_state(
 ) -> None:
     """Test basic climate state reporting."""
     async with run_compiled(yaml_config), api_client_connected() as client:
-        # Get entities and set up state synchronization
-        entities, services = await client.list_entities_services()
-        initial_state_helper = InitialStateHelper(entities)
+        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"
-
-        # Subscribe with the wrapper (no-op callback since we just want initial states)
-        client.subscribe_states(initial_state_helper.on_state_wrapper(lambda _: None))
-
-        # 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")
-
-        # Get the climate entity and its initial state
         test_climate = climate_infos[0]
-        climate_state = initial_state_helper.initial_states.get(test_climate.key)
 
-        assert climate_state is not None, "Climate initial state not found"
-        assert isinstance(climate_state, aioesphomeapi.ClimateState)
-        assert climate_state.mode == ClimateMode.OFF
-        assert climate_state.action == ClimateAction.OFF
-        assert climate_state.current_temperature == 22.0
-        assert climate_state.target_temperature_low == 18.0
-        assert climate_state.target_temperature_high == 24.0
-        assert climate_state.preset == ClimatePreset.HOME
-        assert climate_state.current_humidity == 42.0
-        assert climate_state.target_humidity == 20.0
+        # The thermostat publishes multiple states during setup as the
+        # temperature/humidity sensors come online. Wait for the state to
+        # converge to the expected default values rather than relying on
+        # whichever state happens to arrive first.
+        def is_default_state(state: EntityState) -> bool:
+            return (
+                isinstance(state, ClimateState)
+                and state.key == test_climate.key
+                and state.mode == ClimateMode.OFF
+                and state.action == ClimateAction.OFF
+                and state.current_temperature == 22.0
+                and state.target_temperature_low == 18.0
+                and state.target_temperature_high == 24.0
+                and state.preset == ClimatePreset.HOME
+                and state.current_humidity == 42.0
+                and state.target_humidity == 20.0
+            )
+
+        try:
+            await wait_for_state(client, is_default_state)
+        except TimeoutError:
+            pytest.fail("Climate did not converge to expected default state")