use helper to fix flakey test

tests/integration/state_utils.py

@@ -44,6 +44,7 @@ class InitialStateHelper:
         helper = InitialStateHelper(entities)
         client.subscribe_states(helper.on_state_wrapper(user_callback))
         await helper.wait_for_initial_states()
+        # Access initial states via helper.initial_states[key]
     """
 
     def __init__(self, entities: list[EntityInfo]) -> None:
@@ -63,6 +64,8 @@ class InitialStateHelper:
         self._entities_by_id = {
             (entity.device_id, entity.key): entity for entity in entities
         }
+        # Store initial states by key for test access
+        self.initial_states: dict[int, EntityState] = {}
 
         # Log all entities
         _LOGGER.debug(
@@ -127,6 +130,9 @@ class InitialStateHelper:
 
             # If this entity is waiting for initial state
             if entity_id in self._wait_initial_states:
+                # Store the initial state for test access
+                self.initial_states[state.key] = state
+
                 # Remove from waiting set
                 self._wait_initial_states.discard(entity_id)
 

tests/integration/test_host_mode_climate_basic_state.py

@@ -2,12 +2,11 @@
 
 from __future__ import annotations
 
-import asyncio
-
 import aioesphomeapi
-from aioesphomeapi import ClimateAction, ClimateMode, ClimatePreset, EntityState
+from aioesphomeapi import ClimateAction, ClimateInfo, ClimateMode, ClimatePreset
 import pytest
 
+from .state_utils import InitialStateHelper
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -18,26 +17,27 @@ async def test_host_mode_climate_basic_state(
     api_client_connected: APIClientConnectedFactory,
 ) -> None:
     """Test basic climate state reporting."""
-    loop = asyncio.get_running_loop()
     async with run_compiled(yaml_config), api_client_connected() as client:
-        states: dict[int, EntityState] = {}
+        # Get entities and set up state synchronization
-        climate_future: asyncio.Future[EntityState] = loop.create_future()
+        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"
 
-        def on_state(state: EntityState) -> None:
+        # Subscribe with the wrapper (no-op callback since we just want initial states)
-            states[state.key] = state
+        client.subscribe_states(initial_state_helper.on_state_wrapper(lambda _: None))
-            if (
-                isinstance(state, aioesphomeapi.ClimateState)
-                and not climate_future.done()
-            ):
-                climate_future.set_result(state)
-
-        client.subscribe_states(on_state)
 
+        # Wait for all initial states to be broadcast
         try:
-            climate_state = await asyncio.wait_for(climate_future, timeout=5.0)
+            await initial_state_helper.wait_for_initial_states()
         except TimeoutError:
-            pytest.fail("Climate state not received within 5 seconds")
+            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