[tests] Add require_service, assert_device_name and exact entity lookups (#21150)

tests/integration/README.md

@@ -63,24 +63,45 @@ Consider converting existing integration tests to use `InitialStateHelper` for m
 
 #### Entity Lookup Helpers (`state_utils.py`)
 
-Two helper functions simplify finding entities in test code:
+These helpers simplify finding entities, services and device info in test code:
 
-**`find_entity(entities, object_id_substring, entity_type=None)`**
+**`find_entity(entities, object_id_substring, entity_type=None, *, exact=False)`**
 - Finds an entity by searching for a substring in its `object_id` (case-insensitive)
+- With `exact=True`, the whole `object_id` must match exactly (case-sensitive)
 - Optionally filters by entity type (e.g., `BinarySensorInfo`)
 - Returns `None` if not found
 
-**`require_entity(entities, object_id_substring, entity_type=None, description=None)`**
+**`require_entity(entities, object_id_substring, entity_type=None, description=None, *, exact=False)`**
 - Same as `find_entity` but raises `AssertionError` if not found
 - Use `description` parameter for clearer error messages
 
+**`require_only_entity(entities, entity_type)`**
+- Returns the one entity of `entity_type`; raises `AssertionError` if there is none or more than one
+
+**`require_service(services, name)`**
+- Returns the user service with exactly this name; fails the test with the available names if missing
+
+**`await assert_device_name(client, name)`**
+- Fetches `device_info()`, asserts it is not `None` and has this name, and returns it
+
 ```python
-from aioesphomeapi import BinarySensorInfo
-from .state_utils import require_entity
+from aioesphomeapi import BinarySensorInfo, LockInfo
+from .state_utils import (
+    assert_device_name,
+    require_entity,
+    require_only_entity,
+    require_service,
+)
+
+await assert_device_name(client, "host-test")
+entities, services = await client.list_entities_services()
 
 # Find entities with clear error messages
 binary_sensor = require_entity(entities, "test_sensor", BinarySensorInfo)
 button = require_entity(entities, "set_true", description="Set True button")
+switch = require_entity(entities, "relay", exact=True)  # not "relay_2"
+lock = require_only_entity(entities, LockInfo)
+my_service = require_service(services, "my_service")
 ```
 
 ### Writing Tests
@@ -97,8 +118,7 @@ async def test_my_feature(
     # Write, compile and run the ESPHome device, then connect to API
     async with run_compiled(yaml_config), api_client_connected() as client:
         # Test your feature using the connected client
-        device_info = await client.device_info()
-        assert device_info is not None
+        await assert_device_name(client, "my-test-device")
 ```
 
 ### Creating YAML Fixtures
@@ -260,8 +280,7 @@ except asyncio.TimeoutError:
 ```python
 # Find and execute service
 entities, services = await client.list_entities_services()
-my_service = next((s for s in services if s.name == "my_service"), None)
-assert my_service is not None
+my_service = require_service(services, "my_service")
 
 # Execute with parameters
 await client.execute_service(my_service, {"param1": "value1", "param2": 42})

tests/integration/state_utils.py

@@ -11,10 +11,12 @@ from aioesphomeapi import (
     APIClient,
     BinarySensorState,
     ButtonInfo,
+    DeviceInfo,
     EntityInfo,
     EntityState,
     SensorState,
     TextSensorState,
+    UserService,
 )
 
 _LOGGER = logging.getLogger(__name__)
@@ -113,6 +115,8 @@ def find_entity[T: EntityInfo](
     entities: list[EntityInfo],
     object_id_substring: str,
     entity_type: type[T] | None = None,
+    *,
+    exact: bool = False,
 ) -> T | EntityInfo | None:
     """Find an entity by object_id substring and optionally by type.
 
@@ -120,6 +124,7 @@ def find_entity[T: EntityInfo](
         entities: List of entity info objects from the API
         object_id_substring: Substring to search for in object_id (case-insensitive)
         entity_type: Optional entity type to filter by (e.g., BinarySensorInfo)
+        exact: Match the whole object_id exactly (case-sensitive) instead
 
     Returns:
         The first matching entity, or None if not found
@@ -130,8 +135,12 @@ def find_entity[T: EntityInfo](
     """
     substring_lower = object_id_substring.lower()
     for entity in entities:
-        if substring_lower in entity.object_id.lower() and (
-            entity_type is None or isinstance(entity, entity_type)
+        if entity_type is not None and not isinstance(entity, entity_type):
+            continue
+        if (
+            entity.object_id == object_id_substring
+            if exact
+            else substring_lower in entity.object_id.lower()
         ):
             return entity
     return None
@@ -142,6 +151,8 @@ def require_entity[T: EntityInfo](
     object_id_substring: str,
     entity_type: type[T] | None = None,
     description: str | None = None,
+    *,
+    exact: bool = False,
 ) -> T | EntityInfo:
     """Find an entity or raise AssertionError if not found.
 
@@ -150,6 +161,7 @@ def require_entity[T: EntityInfo](
         object_id_substring: Substring to search for in object_id (case-insensitive)
         entity_type: Optional entity type to filter by (e.g., BinarySensorInfo)
         description: Human-readable description for error message
+        exact: Match the whole object_id exactly (case-sensitive) instead
 
     Returns:
         The first matching entity
@@ -161,14 +173,56 @@ def require_entity[T: EntityInfo](
         binary_sensor = require_entity(entities, "test_sensor", BinarySensorInfo)
         button = require_entity(entities, "set_true", description="Set True button")
     """
-    entity = find_entity(entities, object_id_substring, entity_type)
+    entity = find_entity(entities, object_id_substring, entity_type, exact=exact)
     if entity is None:
-        desc = description or f"entity with '{object_id_substring}' in object_id"
+        desc = description or (
+            f"entity with object_id '{object_id_substring}'"
+            if exact
+            else f"entity with '{object_id_substring}' in object_id"
+        )
         type_info = f" of type {entity_type.__name__}" if entity_type else ""
         raise AssertionError(f"{desc}{type_info} not found in entities")
     return entity
 
 
+def require_only_entity[T: EntityInfo](
+    entities: list[EntityInfo], entity_type: type[T]
+) -> T:
+    """Return the single entity of ``entity_type``.
+
+    Raises AssertionError when there is none or more than one.
+    """
+    matches = [entity for entity in entities if isinstance(entity, entity_type)]
+    assert len(matches) == 1, (
+        f"expected exactly one {entity_type.__name__}, found "
+        f"{[entity.object_id for entity in matches]}"
+    )
+    return matches[0]
+
+
+def require_service(services: list[UserService], name: str) -> UserService:
+    """Return the user service named exactly ``name``.
+
+    Raises AssertionError with the available service names when it is missing.
+    """
+    service = next((s for s in services if s.name == name), None)
+    if service is None:
+        raise AssertionError(
+            f"service '{name}' not found; available: {[s.name for s in services]}"
+        )
+    return service
+
+
+async def assert_device_name(client: APIClient, name: str) -> DeviceInfo:
+    """Fetch device info, assert the device is called ``name`` and return the info."""
+    device_info = await client.device_info()
+    assert device_info is not None, "device_info() returned None"
+    assert device_info.name == name, (
+        f"expected device name '{name}', got '{device_info.name}'"
+    )
+    return device_info
+
+
 def build_key_to_entity_mapping(
     entities: list[EntityInfo], entity_names: list[str]
 ) -> dict[int, str]:

tests/unit_tests/test_integration_state_utils.py

@@ -0,0 +1,77 @@
+"""Tests for the lookup helpers in tests/integration/state_utils.py."""
+
+from __future__ import annotations
+
+from unittest.mock import AsyncMock
+
+from aioesphomeapi import (
+    BinarySensorInfo,
+    ButtonInfo,
+    DeviceInfo,
+    EntityInfo,
+    LightInfo,
+    SensorInfo,
+    UserService,
+)
+import pytest
+
+from tests.integration.state_utils import (
+    assert_device_name,
+    find_entity,
+    require_entity,
+    require_only_entity,
+    require_service,
+)
+
+ENTITIES: list[EntityInfo] = [
+    SensorInfo(object_id="temperature_raw", key=1),
+    SensorInfo(object_id="temperature", key=2),
+    BinarySensorInfo(object_id="Motion", key=3),
+    ButtonInfo(object_id="restart", key=4),
+]
+
+
+def test_find_entity_substring_default() -> None:
+    assert find_entity(ENTITIES, "temperature").key == 1
+    assert find_entity(ENTITIES, "motion").key == 3
+
+
+def test_find_entity_exact() -> None:
+    assert find_entity(ENTITIES, "temperature", exact=True).key == 2
+    assert find_entity(ENTITIES, "temp", exact=True) is None
+    assert find_entity(ENTITIES, "motion", exact=True) is None
+    assert find_entity(ENTITIES, "temperature", ButtonInfo, exact=True) is None
+
+
+def test_require_entity_exact() -> None:
+    assert require_entity(ENTITIES, "Motion", BinarySensorInfo, exact=True).key == 3
+    with pytest.raises(AssertionError, match="object_id 'temp' of type SensorInfo"):
+        require_entity(ENTITIES, "temp", SensorInfo, exact=True)
+
+
+def test_require_only_entity() -> None:
+    assert require_only_entity(ENTITIES, ButtonInfo).key == 4
+    with pytest.raises(AssertionError, match="exactly one SensorInfo"):
+        require_only_entity(ENTITIES, SensorInfo)
+    with pytest.raises(AssertionError, match=r"exactly one LightInfo, found \[\]"):
+        require_only_entity(ENTITIES, LightInfo)
+
+
+def test_require_service() -> None:
+    services = [UserService(name="set_level", key=1), UserService(name="set", key=2)]
+    assert require_service(services, "set").key == 2
+    with pytest.raises(AssertionError, match=r"'missing' not found.*set_level"):
+        require_service(services, "missing")
+
+
+@pytest.mark.asyncio
+async def test_assert_device_name() -> None:
+    info = DeviceInfo(name="host-test")
+    client = AsyncMock()
+    client.device_info.return_value = info
+    assert await assert_device_name(client, "host-test") is info
+    with pytest.raises(AssertionError, match="expected device name 'other'"):
+        await assert_device_name(client, "other")
+    client.device_info.return_value = None
+    with pytest.raises(AssertionError, match="returned None"):
+        await assert_device_name(client, "host-test")