fix flakey

tests/integration/README.md

@@ -8,7 +8,7 @@ This directory contains end-to-end integration tests for ESPHome, focusing on te
 - `const.py` - Constants used throughout the integration tests
 - `types.py` - Type definitions for fixtures and functions
 - `state_utils.py` - State handling utilities (e.g., `InitialStateHelper`)
-- `sensor_test_utils.py` - Sensor-specific test utilities
+- `sensor_utils.py` - Sensor-specific test utilities
 - `fixtures/` - YAML configuration files for tests
 - `test_*.py` - Individual test files
 

tests/integration/state_utils.py

@@ -10,6 +10,28 @@ from aioesphomeapi import ButtonInfo, EntityInfo, EntityState
 _LOGGER = logging.getLogger(__name__)
 
 
+def build_key_to_entity_mapping(
+    entities: list[EntityInfo], entity_names: list[str]
+) -> dict[int, str]:
+    """Build a mapping from entity keys to entity names.
+
+    Args:
+        entities: List of entity info objects from the API
+        entity_names: List of entity names to search for in object_ids
+
+    Returns:
+        Dictionary mapping entity keys to entity names
+    """
+    key_to_entity: dict[int, str] = {}
+    for entity in entities:
+        obj_id = entity.object_id.lower()
+        for entity_name in entity_names:
+            if entity_name in obj_id:
+                key_to_entity[entity.key] = entity_name
+                break
+    return key_to_entity
+
+
 class InitialStateHelper:
     """Helper to wait for initial states before processing test states.
 

tests/integration/test_sensor_filters_ring_buffer.py

@@ -7,8 +7,7 @@ import asyncio
 from aioesphomeapi import EntityState, SensorState
 import pytest
 
-from .sensor_test_utils import build_key_to_sensor_mapping
+from .state_utils import InitialStateHelper, build_key_to_entity_mapping
-from .state_utils import InitialStateHelper
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -67,7 +66,7 @@ async def test_sensor_filters_ring_buffer(
         entities, services = await client.list_entities_services()
 
         # Build key-to-sensor mapping
-        key_to_sensor = build_key_to_sensor_mapping(
+        key_to_sensor = build_key_to_entity_mapping(
             entities,
             [
                 "sliding_min",

tests/integration/test_sensor_filters_sliding_window.py

@@ -7,8 +7,7 @@ import asyncio
 from aioesphomeapi import EntityState, SensorState
 import pytest
 
-from .sensor_test_utils import build_key_to_sensor_mapping
+from .state_utils import InitialStateHelper, build_key_to_entity_mapping
-from .state_utils import InitialStateHelper
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
@@ -98,7 +97,7 @@ async def test_sensor_filters_sliding_window(
         entities, services = await client.list_entities_services()
 
         # Build key-to-sensor mapping
-        key_to_sensor = build_key_to_sensor_mapping(
+        key_to_sensor = build_key_to_entity_mapping(
             entities,
             [
                 "min_sensor",
@@ -245,7 +244,7 @@ async def test_sensor_filters_nan_handling(
         entities, services = await client.list_entities_services()
 
         # Build key-to-sensor mapping
-        key_to_sensor = build_key_to_sensor_mapping(entities, ["min_nan", "max_nan"])
+        key_to_sensor = build_key_to_entity_mapping(entities, ["min_nan", "max_nan"])
 
         # Set up initial state helper with all entities
         initial_state_helper = InitialStateHelper(entities)
@@ -345,7 +344,7 @@ async def test_sensor_filters_ring_buffer_wraparound(
         entities, services = await client.list_entities_services()
 
         # Build key-to-sensor mapping
-        key_to_sensor = build_key_to_sensor_mapping(entities, ["wraparound_min"])
+        key_to_sensor = build_key_to_entity_mapping(entities, ["wraparound_min"])
 
         # Set up initial state helper with all entities
         initial_state_helper = InitialStateHelper(entities)