From 26016b29f7aaa130a1e681e94ff1e4102a1ef811 Mon Sep 17 00:00:00 2001
From: dougiteixeira <31328123+dougiteixeira@users.noreply.github.com>
Date: Mon, 26 Jun 2023 13:05:11 -0300
Subject: [PATCH] Add the device of the source entity in the helper entities
 for Threshold (#94753)

---
 .../components/threshold/binary_sensor.py     | 41 +++++++++++++++-
 .../threshold/test_binary_sensor.py           | 47 +++++++++++++++++++
 2 files changed, 86 insertions(+), 2 deletions(-)

diff --git a/homeassistant/components/threshold/binary_sensor.py b/homeassistant/components/threshold/binary_sensor.py
index 538655ec0ce3b6..f7b8c9c097c7b9 100644
--- a/homeassistant/components/threshold/binary_sensor.py
+++ b/homeassistant/components/threshold/binary_sensor.py
@@ -22,7 +22,12 @@
     STATE_UNKNOWN,
 )
 from homeassistant.core import Event, HomeAssistant, callback
-from homeassistant.helpers import config_validation as cv, entity_registry as er
+from homeassistant.helpers import (
+    config_validation as cv,
+    device_registry as dr,
+    entity_registry as er,
+)
+from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.event import async_track_state_change_event
 from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
@@ -73,6 +78,28 @@ async def async_setup_entry(
     entity_id = er.async_validate_entity_id(
         registry, config_entry.options[CONF_ENTITY_ID]
     )
+
+    source_entity = registry.async_get(entity_id)
+    dev_reg = dr.async_get(hass)
+    # Resolve source entity device
+    if (
+        (source_entity is not None)
+        and (source_entity.device_id is not None)
+        and (
+            (
+                device := dev_reg.async_get(
+                    device_id=source_entity.device_id,
+                )
+            )
+            is not None
+        )
+    ):
+        device_info = DeviceInfo(
+            identifiers=device.identifiers,
+        )
+    else:
+        device_info = None
+
     hysteresis = config_entry.options[CONF_HYSTERESIS]
     lower = config_entry.options[CONF_LOWER]
     name = config_entry.title
@@ -82,7 +109,15 @@ async def async_setup_entry(
     async_add_entities(
         [
             ThresholdSensor(
-                hass, entity_id, name, lower, upper, hysteresis, device_class, unique_id
+                hass,
+                entity_id,
+                name,
+                lower,
+                upper,
+                hysteresis,
+                device_class,
+                unique_id,
+                device_info=device_info,
             )
         ]
     )
@@ -138,9 +173,11 @@ def __init__(
         hysteresis: float,
         device_class: BinarySensorDeviceClass | None,
         unique_id: str | None,
+        device_info: DeviceInfo | None = None,
     ) -> None:
         """Initialize the Threshold sensor."""
         self._attr_unique_id = unique_id
+        self._attr_device_info = device_info
         self._entity_id = entity_id
         self._name = name
         if lower is not None:
diff --git a/tests/components/threshold/test_binary_sensor.py b/tests/components/threshold/test_binary_sensor.py
index 9e11195d878f11..2180d0aed7f278 100644
--- a/tests/components/threshold/test_binary_sensor.py
+++ b/tests/components/threshold/test_binary_sensor.py
@@ -2,6 +2,7 @@
 
 import pytest
 
+from homeassistant.components.threshold.const import DOMAIN
 from homeassistant.const import (
     ATTR_UNIT_OF_MEASUREMENT,
     STATE_UNAVAILABLE,
@@ -9,8 +10,11 @@
     UnitOfTemperature,
 )
 from homeassistant.core import HomeAssistant
+from homeassistant.helpers import device_registry as dr, entity_registry as er
 from homeassistant.setup import async_setup_component
 
+from tests.common import MockConfigEntry
+
 
 async def test_sensor_upper(hass: HomeAssistant) -> None:
     """Test if source is above threshold."""
@@ -585,3 +589,46 @@ async def test_sensor_no_lower_upper(
     await hass.async_block_till_done()
 
     assert "Lower or Upper thresholds not provided" in caplog.text
+
+
+async def test_device_id(hass: HomeAssistant) -> None:
+    """Test for source entity device for Threshold."""
+    device_registry = dr.async_get(hass)
+    entity_registry = er.async_get(hass)
+
+    source_config_entry = MockConfigEntry()
+    source_device_entry = device_registry.async_get_or_create(
+        config_entry_id=source_config_entry.entry_id,
+        identifiers={("sensor", "identifier_test")},
+    )
+    source_entity = entity_registry.async_get_or_create(
+        "sensor",
+        "test",
+        "source",
+        config_entry=source_config_entry,
+        device_id=source_device_entry.id,
+    )
+    await hass.async_block_till_done()
+    assert entity_registry.async_get("sensor.test_source") is not None
+
+    utility_meter_config_entry = MockConfigEntry(
+        data={},
+        domain=DOMAIN,
+        options={
+            "entity_id": "sensor.test_source",
+            "hysteresis": 0.0,
+            "lower": -2.0,
+            "name": "Threshold",
+            "upper": None,
+        },
+        title="Threshold",
+    )
+
+    utility_meter_config_entry.add_to_hass(hass)
+
+    assert await hass.config_entries.async_setup(utility_meter_config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    utility_meter_entity = entity_registry.async_get("binary_sensor.threshold")
+    assert utility_meter_entity is not None
+    assert utility_meter_entity.device_id == source_entity.device_id