[cover] Use InitialStateHelper in integration test

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