Add common template test framework to vacuum platform (home-assistant#157846)

Author: Petro31

tests/components/template/conftest.py

@@ -1,5 +1,6 @@
 """template conftest."""
 
+from dataclasses import dataclass
 from enum import Enum
 
 import pytest
@@ -8,6 +9,7 @@
 from homeassistant.config_entries import SOURCE_USER
 from homeassistant.core import HomeAssistant, ServiceCall
 from homeassistant.data_entry_flow import FlowResultType
+from homeassistant.helpers import entity_registry as er
 from homeassistant.helpers.typing import ConfigType
 from homeassistant.setup import async_setup_component
 
@@ -40,19 +42,37 @@ def make_test_trigger(*entities: str) -> dict:
     }
 
 
+async def async_trigger(
+    hass: HomeAssistant, entity_id: str, state: str | None = None
+) -> None:
+    """Trigger a state change."""
+    hass.states.async_set(entity_id, state)
+    await hass.async_block_till_done()
+
+
 async def async_setup_legacy_platforms(
     hass: HomeAssistant,
     domain: str,
-    slug: str,
+    slug: str | None,
     count: int,
-    config: ConfigType,
+    config: ConfigType | list[ConfigType],
 ) -> None:
     """Do setup of any legacy platform that supports a keyed dictionary of template entities."""
+    if slug is None:
+        # Lock and Weather platforms do not use a slug
+        if isinstance(config, list):
+            config = {domain: [{"platform": "template", **item} for item in config]}
+        else:
+            config = {domain: {"platform": "template", **config}}
+    else:
+        assert isinstance(config, dict)
+        config = {domain: {"platform": "template", slug: config}}
+
     with assert_setup_component(count, domain):
         assert await async_setup_component(
             hass,
             domain,
-            {domain: {"platform": "template", slug: config}},
+            config,
         )
 
     await hass.async_block_till_done()
@@ -64,16 +84,15 @@ async def async_setup_modern_state_format(
     hass: HomeAssistant,
     domain: str,
     count: int,
-    config: ConfigType,
-    extra_config: ConfigType | None = None,
+    config: ConfigType | list[ConfigType],
+    extra_section_config: ConfigType | None = None,
 ) -> None:
     """Do setup of template integration via modern format."""
-    extra = extra_config or {}
     with assert_setup_component(count, template.DOMAIN):
         assert await async_setup_component(
             hass,
             template.DOMAIN,
-            {"template": {domain: config, **extra}},
+            {"template": {domain: config, **(extra_section_config or {})}},
         )
 
     await hass.async_block_till_done()
@@ -86,12 +105,11 @@ async def async_setup_modern_trigger_format(
     domain: str,
     trigger: dict,
     count: int,
-    config: ConfigType,
-    extra_config: ConfigType | None = None,
+    config: ConfigType | list[ConfigType],
+    extra_section_config: ConfigType | None = None,
 ) -> None:
     """Do setup of template integration via trigger format."""
-    extra = extra_config or {}
-    config = {"template": {domain: config, **trigger, **extra}}
+    config = {"template": {domain: config, **trigger, **(extra_section_config or {})}}
 
     with assert_setup_component(count, template.DOMAIN):
         assert await async_setup_component(
@@ -105,6 +123,164 @@ async def async_setup_modern_trigger_format(
     await hass.async_block_till_done()
 
 
+@dataclass(frozen=True)
+class TemplatePlatformSetup:
+    """Template Platform Setup Information."""
+
+    domain: str
+    legacy_slug: str | None
+    object_id: str
+    trigger: ConfigType
+
+    @property
+    def entity_id(self) -> str:
+        """Return test entity ID."""
+        return f"{self.domain}.{self.object_id}"
+
+
+async def setup_entity(
+    hass: HomeAssistant,
+    platform_setup: TemplatePlatformSetup,
+    style: ConfigurationStyle,
+    count: int,
+    config: ConfigType,
+    state_template: str | None = None,
+    extra_config: ConfigType | None = None,
+    attributes: ConfigType | None = None,
+    extra_section_config: ConfigType | None = None,
+) -> None:
+    """Do setup of a template entity based on the configuration style."""
+    if style == ConfigurationStyle.LEGACY:
+        await async_setup_legacy_platforms(
+            hass,
+            platform_setup.domain,
+            platform_setup.legacy_slug,
+            count,
+            {
+                platform_setup.object_id: {
+                    **({"value_template": state_template} if state_template else {}),
+                    **config,
+                    **(extra_config or {}),
+                    **({"attribute_templates": attributes} if attributes else {}),
+                }
+            },
+        )
+        return
+
+    entity_config = {
+        "name": platform_setup.object_id,
+        **({"state": state_template} if state_template else {}),
+        **config,
+        **({"attributes": attributes} if attributes else {}),
+        **(extra_config or {}),
+    }
+    if style == ConfigurationStyle.MODERN:
+        await async_setup_modern_state_format(
+            hass, platform_setup.domain, count, entity_config, extra_section_config
+        )
+    elif style == ConfigurationStyle.TRIGGER:
+        await async_setup_modern_trigger_format(
+            hass,
+            platform_setup.domain,
+            platform_setup.trigger,
+            count,
+            entity_config,
+            extra_section_config,
+        )
+
+
+async def setup_and_test_unique_id(
+    hass: HomeAssistant,
+    platform_setup: TemplatePlatformSetup,
+    style: ConfigurationStyle,
+    entity_config: ConfigType | None,
+) -> None:
+    """Setup 2 entities with the same unique_id and verify only 1 entity is created.
+
+    The entity_config not provide name or unique_id, those are added automatically.
+    """
+    entity_config = {"unique_id": "not-so_-unique-anymore", **(entity_config or {})}
+    if style == ConfigurationStyle.LEGACY:
+        if platform_setup.legacy_slug is None:
+            config = [
+                {"name": "template_entity_1", **entity_config},
+                {"name": "template_entity_2", **entity_config},
+            ]
+        else:
+            config = {
+                "template_entity_1": entity_config,
+                "template_entity_2": entity_config,
+            }
+        await async_setup_legacy_platforms(
+            hass, platform_setup.domain, platform_setup.legacy_slug, 1, config
+        )
+    elif style == ConfigurationStyle.MODERN:
+        await async_setup_modern_state_format(
+            hass,
+            platform_setup.domain,
+            1,
+            [
+                {"name": "template_entity_1", **entity_config},
+                {"name": "template_entity_2", **entity_config},
+            ],
+        )
+    elif style == ConfigurationStyle.TRIGGER:
+        await async_setup_modern_trigger_format(
+            hass,
+            platform_setup.domain,
+            platform_setup.trigger,
+            1,
+            [
+                {"name": "template_entity_1", **entity_config},
+                {"name": "template_entity_2", **entity_config},
+            ],
+        )
+
+    assert len(hass.states.async_all(platform_setup.domain)) == 1
+
+
+async def setup_and_test_nested_unique_id(
+    hass: HomeAssistant,
+    platform_setup: TemplatePlatformSetup,
+    style: ConfigurationStyle,
+    entity_registry: er.EntityRegistry,
+    entity_config: ConfigType | None,
+) -> None:
+    """Setup 2 entities with unique unique_ids in a template section that contains a unique_id.
+
+    The test will verify that 2 entities are created where the unique_id appends the
+    section unique_id to each entity unique_id.
+
+    The entity_config should not provide name or unique_id, those are added automatically.
+    """
+    entities = [
+        {"name": "test_a", "unique_id": "a", **(entity_config or {})},
+        {"name": "test_b", "unique_id": "b", **(entity_config or {})},
+    ]
+    extra_section_config = {"unique_id": "x"}
+    if style == ConfigurationStyle.MODERN:
+        await async_setup_modern_state_format(
+            hass, platform_setup.domain, 1, entities, extra_section_config
+        )
+    elif style == ConfigurationStyle.TRIGGER:
+        await async_setup_modern_trigger_format(
+            hass,
+            platform_setup.domain,
+            platform_setup.trigger,
+            1,
+            entities,
+            extra_section_config,
+        )
+
+    assert len(hass.states.async_all(platform_setup.domain)) == 2
+
+    entry = entity_registry.async_get(f"{platform_setup.domain}.test_a")
+    assert entry.unique_id == "x-a"
+
+    entry = entity_registry.async_get(f"{platform_setup.domain}.test_b")
+    assert entry.unique_id == "x-b"
+
+
 @pytest.fixture
 def calls(hass: HomeAssistant) -> list[ServiceCall]:
     """Track calls to a mock service."""