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

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")