From 206ea9d237a760e42d8e00f03ab8021894fd0727 Mon Sep 17 00:00:00 2001
From: Paulus Schoutsen <balloob@gmail.com>
Date: Wed, 30 Mar 2022 01:11:09 -0700
Subject: [PATCH] Expand group lights/covers etc (#68875)

Co-authored-by: Franck Nijhof <git@frenck.dev>
---
 homeassistant/helpers/template.py |  8 +++++++-
 tests/helpers/test_template.py    | 26 ++++++++++++++++++++++++++
 2 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/homeassistant/helpers/template.py b/homeassistant/helpers/template.py
index 2b93b69fe37..0d849f67d9a 100644
--- a/homeassistant/helpers/template.py
+++ b/homeassistant/helpers/template.py
@@ -887,6 +887,9 @@ def result_as_boolean(template_result: Any | None) -> bool:
 
 def expand(hass: HomeAssistant, *args: Any) -> Iterable[State]:
     """Expand out any groups into entity states."""
+    # circular import.
+    from . import entity as entity_helper  # pylint: disable=import-outside-toplevel
+
     search = list(args)
     found = {}
     while search:
@@ -904,7 +907,10 @@ def expand(hass: HomeAssistant, *args: Any) -> Iterable[State]:
             # ignore other types
             continue
 
-        if entity_id.startswith(_GROUP_DOMAIN_PREFIX):
+        if entity_id.startswith(_GROUP_DOMAIN_PREFIX) or (
+            (source := entity_helper.entity_sources(hass).get(entity_id))
+            and source["domain"] == "group"
+        ):
             # Collect state will be called in here since it's wrapped
             group_entities = entity.attributes.get(ATTR_ENTITY_ID)
             if group_entities:
diff --git a/tests/helpers/test_template.py b/tests/helpers/test_template.py
index bbb5e11be7f..7b0ceb313af 100644
--- a/tests/helpers/test_template.py
+++ b/tests/helpers/test_template.py
@@ -2079,6 +2079,32 @@ async def test_expand(hass):
     )
     assert info.rate_limit is None
 
+    # With group entities
+    hass.states.async_set("light.first", "on")
+    hass.states.async_set("light.second", "off")
+
+    assert await async_setup_component(
+        hass,
+        "light",
+        {
+            "light": {
+                "platform": "group",
+                "name": "Grouped",
+                "entities": ["light.first", "light.second"],
+            }
+        },
+    )
+    await hass.async_block_till_done()
+
+    info = render_to_info(
+        hass, "{{ expand('light.grouped') | map(attribute='entity_id') | join(', ') }}"
+    )
+    assert_result_info(
+        info,
+        "light.first, light.second",
+        ["light.grouped", "light.first", "light.second"],
+    )
+
 
 async def test_device_entities(hass):
     """Test device_entities function."""