diff --git a/tests/integration/test_cover_control_action.py b/tests/integration/test_cover_control_action.py index fee89555fa3..206079ca89a 100644 --- a/tests/integration/test_cover_control_action.py +++ b/tests/integration/test_cover_control_action.py @@ -5,11 +5,14 @@ work correctly with the per-instance bitmask field storage. Exercises multiple field combinations to cover the bitmask variants. """ -import asyncio -from typing import Any +from __future__ import annotations +import asyncio + +from aioesphomeapi import ButtonInfo, CoverInfo, CoverState, EntityState import pytest +from .state_utils import InitialStateHelper, require_entity from .types import APIClientConnectedFactory, RunCompiledFunction @@ -20,31 +23,37 @@ async def test_cover_control_action( api_client_connected: APIClientConnectedFactory, ) -> None: """Test cover ControlAction/CoverPublishAction with constants and lambdas.""" + loop = asyncio.get_running_loop() async with run_compiled(yaml_config), api_client_connected() as client: - state_futures: dict[int, asyncio.Future[Any]] = {} + cover_state_future: asyncio.Future[CoverState] | None = None - def on_state(state: Any) -> None: - if state.key in state_futures and not state_futures[state.key].done(): - state_futures[state.key].set_result(state) + def on_state(state: EntityState) -> None: + if ( + isinstance(state, CoverState) + and cover_state_future is not None + and not cover_state_future.done() + ): + cover_state_future.set_result(state) - client.subscribe_states(on_state) - - entities = await client.list_entities_services() - cover = next(e for e in entities[0] if e.object_id == "test_cover") - buttons = {e.name: e for e in entities[0] if hasattr(e, "name")} - - async def wait_for_state(key: int, timeout: float = 5.0) -> Any: - loop = asyncio.get_running_loop() - state_futures[key] = loop.create_future() + async def wait_for_cover_state(timeout: float = 5.0) -> CoverState: + nonlocal cover_state_future + cover_state_future = loop.create_future() try: - return await asyncio.wait_for(state_futures[key], timeout) + return await asyncio.wait_for(cover_state_future, timeout) finally: - state_futures.pop(key, None) + cover_state_future = None - async def press_and_wait(button_name: str) -> Any: - btn = buttons[button_name] + entities, _ = await client.list_entities_services() + initial_state_helper = InitialStateHelper(entities) + client.subscribe_states(initial_state_helper.on_state_wrapper(on_state)) + await initial_state_helper.wait_for_initial_states() + + require_entity(entities, "test_cover", CoverInfo) + + async def press_and_wait(name: str) -> CoverState: + btn = require_entity(entities, name.lower().replace(" ", "_"), ButtonInfo) client.button_command(btn.key) - return await wait_for_state(cover.key) + return await wait_for_cover_state() # Test 1: position only (mask 2) state = await press_and_wait("Set Position") @@ -59,26 +68,19 @@ async def test_cover_control_action( assert state.position == pytest.approx(0.25, abs=0.01) assert state.tilt == pytest.approx(0.30, abs=0.01) - # Test 4: stop (mask 1) — stop on a non-moving cover should not change state - # We just verify the action runs without error; no state change expected. - btn = buttons["Stop Cover"] - client.button_command(btn.key) - # Give the action a moment to execute - await asyncio.sleep(0.1) - - # Test 5: state: OPEN (CONF_STATE alias for position 1.0) + # Test 4: state: OPEN (CONF_STATE alias for position 1.0) state = await press_and_wait("Open State") assert state.position == pytest.approx(1.0, abs=0.01) - # Test 6: lambda position + # Test 5: lambda position (test_position global = 0.42) state = await press_and_wait("Lambda Position") assert state.position == pytest.approx(0.42, abs=0.01) - # Test 7: cover.template.publish position only + # Test 6: cover.template.publish position only state = await press_and_wait("Publish Pos") assert state.position == pytest.approx(0.6, abs=0.01) - # Test 8: cover.template.publish current_operation only + # Test 7: cover.template.publish current_operation only state = await press_and_wait("Publish Op") - # Operation 1 = OPENING + # CoverOperation.OPENING == 1 assert state.current_operation == 1