From 9d4ad1821e2a75d3b8848a55e6d4106c3da57737 Mon Sep 17 00:00:00 2001
From: Tobias Sauerwein <cgtobi@users.noreply.github.com>
Date: Wed, 14 Apr 2021 14:12:26 +0200
Subject: [PATCH] Fix logic of entity id extraction (#49164)

Co-authored-by: Paulus Schoutsen <paulus@home-assistant.io>
---
 homeassistant/helpers/service.py | 12 +++++++--
 tests/helpers/test_service.py    | 43 +++++++++++++++++++++++++++++++-
 2 files changed, 52 insertions(+), 3 deletions(-)

diff --git a/homeassistant/helpers/service.py b/homeassistant/helpers/service.py
index 4e484c6aaab..9dec919d4b5 100644
--- a/homeassistant/helpers/service.py
+++ b/homeassistant/helpers/service.py
@@ -362,8 +362,16 @@ async def async_extract_referenced_entity_ids(
         return selected
 
     for ent_entry in ent_reg.entities.values():
-        if ent_entry.area_id in selector.area_ids or (
-            not ent_entry.area_id and ent_entry.device_id in selected.referenced_devices
+        if (
+            # when area matches the target area
+            ent_entry.area_id in selector.area_ids
+            # when device matches a referenced devices with no explicitly set area
+            or (
+                not ent_entry.area_id
+                and ent_entry.device_id in selected.referenced_devices
+            )
+            # when device matches target device
+            or ent_entry.device_id in selector.device_ids
         ):
             selected.indirectly_referenced.add(ent_entry.entity_id)
 
diff --git a/tests/helpers/test_service.py b/tests/helpers/test_service.py
index 7538c0f6f2c..7e18547145b 100644
--- a/tests/helpers/test_service.py
+++ b/tests/helpers/test_service.py
@@ -95,6 +95,7 @@ def area_mock(hass):
     device_in_area = dev_reg.DeviceEntry(area_id="test-area")
     device_no_area = dev_reg.DeviceEntry(id="device-no-area-id")
     device_diff_area = dev_reg.DeviceEntry(area_id="diff-area")
+    device_area_a = dev_reg.DeviceEntry(id="device-area-a-id", area_id="area-a")
 
     mock_device_registry(
         hass,
@@ -102,6 +103,7 @@ def area_mock(hass):
             device_in_area.id: device_in_area,
             device_no_area.id: device_no_area,
             device_diff_area.id: device_diff_area,
+            device_area_a.id: device_area_a,
         },
     )
 
@@ -119,7 +121,7 @@ def area_mock(hass):
     )
     entity_in_other_area = ent_reg.RegistryEntry(
         entity_id="light.in_other_area",
-        unique_id="in-other-area-id",
+        unique_id="in-area-a-id",
         platform="test",
         device_id=device_in_area.id,
         area_id="other-area",
@@ -143,6 +145,20 @@ def area_mock(hass):
         platform="test",
         device_id=device_diff_area.id,
     )
+    entity_in_area_a = ent_reg.RegistryEntry(
+        entity_id="light.in_area_a",
+        unique_id="in-area-a-id",
+        platform="test",
+        device_id=device_area_a.id,
+        area_id="area-a",
+    )
+    entity_in_area_b = ent_reg.RegistryEntry(
+        entity_id="light.in_area_b",
+        unique_id="in-area-b-id",
+        platform="test",
+        device_id=device_area_a.id,
+        area_id="area-b",
+    )
     mock_registry(
         hass,
         {
@@ -152,6 +168,8 @@ def area_mock(hass):
             entity_assigned_to_area.entity_id: entity_assigned_to_area,
             entity_no_area.entity_id: entity_no_area,
             entity_diff_area.entity_id: entity_diff_area,
+            entity_in_area_a.entity_id: entity_in_area_a,
+            entity_in_area_b.entity_id: entity_in_area_b,
         },
     )
 
@@ -399,6 +417,29 @@ async def test_extract_entity_ids_from_area(hass, area_mock):
     )
 
 
+async def test_extract_entity_ids_from_devices(hass, area_mock):
+    """Test extract_entity_ids method with devices."""
+    assert await service.async_extract_entity_ids(
+        hass, ha.ServiceCall("light", "turn_on", {"device_id": "device-no-area-id"})
+    ) == {
+        "light.no_area",
+    }
+
+    assert await service.async_extract_entity_ids(
+        hass, ha.ServiceCall("light", "turn_on", {"device_id": "device-area-a-id"})
+    ) == {
+        "light.in_area_a",
+        "light.in_area_b",
+    }
+
+    assert (
+        await service.async_extract_entity_ids(
+            hass, ha.ServiceCall("light", "turn_on", {"device_id": "non-existing-id"})
+        )
+        == set()
+    )
+
+
 async def test_async_get_all_descriptions(hass):
     """Test async_get_all_descriptions."""
     group = hass.components.group