From 1d1b5ab3451c0001ca8d64598ed8f7b54fa359a3 Mon Sep 17 00:00:00 2001
From: Raman Gupta <7243222+raman325@users.noreply.github.com>
Date: Mon, 30 Aug 2021 16:09:41 -0400
Subject: [PATCH] Fix area_id and area_name template functions (#55470)

---
 homeassistant/helpers/template.py | 23 ++++++++++++++++++-----
 tests/helpers/test_template.py    | 20 ++++++++++++++++++++
 2 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/homeassistant/helpers/template.py b/homeassistant/helpers/template.py
index a831e8d156d..ade580694c8 100644
--- a/homeassistant/helpers/template.py
+++ b/homeassistant/helpers/template.py
@@ -957,6 +957,7 @@ def area_id(hass: HomeAssistant, lookup_value: str) -> str | None:
         return area.id
 
     ent_reg = entity_registry.async_get(hass)
+    dev_reg = device_registry.async_get(hass)
     # Import here, not at top-level to avoid circular import
     from homeassistant.helpers import (  # pylint: disable=import-outside-toplevel
         config_validation as cv,
@@ -968,10 +969,14 @@ def area_id(hass: HomeAssistant, lookup_value: str) -> str | None:
         pass
     else:
         if entity := ent_reg.async_get(lookup_value):
-            return entity.area_id
+            # If entity has an area ID, return that
+            if entity.area_id:
+                return entity.area_id
+            # If entity has a device ID, return the area ID for the device
+            if entity.device_id and (device := dev_reg.async_get(entity.device_id)):
+                return device.area_id
 
-    # Check if this could be a device ID (hex string)
-    dev_reg = device_registry.async_get(hass)
+    # Check if this could be a device ID
     if device := dev_reg.async_get(lookup_value):
         return device.area_id
 
@@ -992,6 +997,7 @@ def area_name(hass: HomeAssistant, lookup_value: str) -> str | None:
     if area:
         return area.name
 
+    dev_reg = device_registry.async_get(hass)
     ent_reg = entity_registry.async_get(hass)
     # Import here, not at top-level to avoid circular import
     from homeassistant.helpers import (  # pylint: disable=import-outside-toplevel
@@ -1004,11 +1010,18 @@ def area_name(hass: HomeAssistant, lookup_value: str) -> str | None:
         pass
     else:
         if entity := ent_reg.async_get(lookup_value):
+            # If entity has an area ID, get the area name for that
             if entity.area_id:
                 return _get_area_name(area_reg, entity.area_id)
-            return None
+            # If entity has a device ID and the device exists with an area ID, get the
+            # area name for that
+            if (
+                entity.device_id
+                and (device := dev_reg.async_get(entity.device_id))
+                and device.area_id
+            ):
+                return _get_area_name(area_reg, device.area_id)
 
-    dev_reg = device_registry.async_get(hass)
     if (device := dev_reg.async_get(lookup_value)) and device.area_id:
         return _get_area_name(area_reg, device.area_id)
 
diff --git a/tests/helpers/test_template.py b/tests/helpers/test_template.py
index 7a2776fd5b2..64b075b685a 100644
--- a/tests/helpers/test_template.py
+++ b/tests/helpers/test_template.py
@@ -1828,6 +1828,16 @@ async def test_area_id(hass):
     assert_result_info(info, area_entry_entity_id.id)
     assert info.rate_limit is None
 
+    # Make sure that when entity doesn't have an area but its device does, that's what
+    # gets returned
+    entity_entry = entity_registry.async_update_entity(
+        entity_entry.entity_id, area_id=area_entry_entity_id.id
+    )
+
+    info = render_to_info(hass, f"{{{{ area_id('{entity_entry.entity_id}') }}}}")
+    assert_result_info(info, area_entry_entity_id.id)
+    assert info.rate_limit is None
+
 
 async def test_area_name(hass):
     """Test area_name function."""
@@ -1897,6 +1907,16 @@ async def test_area_name(hass):
     assert_result_info(info, area_entry.name)
     assert info.rate_limit is None
 
+    # Make sure that when entity doesn't have an area but its device does, that's what
+    # gets returned
+    entity_entry = entity_registry.async_update_entity(
+        entity_entry.entity_id, area_id=None
+    )
+
+    info = render_to_info(hass, f"{{{{ area_name('{entity_entry.entity_id}') }}}}")
+    assert_result_info(info, area_entry.name)
+    assert info.rate_limit is None
+
 
 def test_closest_function_to_coord(hass):
     """Test closest function to coord."""