diff --git a/tests/integration/state_utils.py b/tests/integration/state_utils.py index 58d6d2790f..6434a41ddf 100644 --- a/tests/integration/state_utils.py +++ b/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) diff --git a/tests/integration/test_host_mode_climate_basic_state.py b/tests/integration/test_host_mode_climate_basic_state.py index 4697342a99..7d871ed5a8 100644 --- a/tests/integration/test_host_mode_climate_basic_state.py +++ b/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] = {} - climate_future: asyncio.Future[EntityState] = loop.create_future() + # Get entities and set up state synchronization + 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: - states[state.key] = state - if ( - isinstance(state, aioesphomeapi.ClimateState) - and not climate_future.done() - ): - climate_future.set_result(state) - - client.subscribe_states(on_state) + # 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: - 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