From e5afff7f989e686631127bc85854d4c791103e7a Mon Sep 17 00:00:00 2001
From: dougiteixeira <31328123+dougiteixeira@users.noreply.github.com>
Date: Thu, 22 Jun 2023 19:04:51 -0300
Subject: [PATCH] Add the device of the source entity in the helper entities
 for Riemann sum integral (#94727)

---
 .../components/integration/sensor.py          | 31 ++++++++++-
 tests/components/integration/test_sensor.py   | 52 ++++++++++++++++++-
 2 files changed, 81 insertions(+), 2 deletions(-)

diff --git a/homeassistant/components/integration/sensor.py b/homeassistant/components/integration/sensor.py
index b28b426d3af..ad0f96dd540 100644
--- a/homeassistant/components/integration/sensor.py
+++ b/homeassistant/components/integration/sensor.py
@@ -28,7 +28,12 @@ from homeassistant.const import (
     UnitOfTime,
 )
 from homeassistant.core import Event, HomeAssistant, State, 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
@@ -140,6 +145,27 @@ async def async_setup_entry(
         registry, config_entry.options[CONF_SOURCE_SENSOR]
     )
 
+    source_entity = er.EntityRegistry.async_get(registry, source_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
+
     unit_prefix = config_entry.options[CONF_UNIT_PREFIX]
     if unit_prefix == "none":
         unit_prefix = None
@@ -152,6 +178,7 @@ async def async_setup_entry(
         unique_id=config_entry.entry_id,
         unit_prefix=unit_prefix,
         unit_time=config_entry.options[CONF_UNIT_TIME],
+        device_info=device_info,
     )
 
     async_add_entities([integral])
@@ -194,6 +221,7 @@ class IntegrationSensor(RestoreSensor):
         unique_id: str | None,
         unit_prefix: str | None,
         unit_time: UnitOfTime,
+        device_info: DeviceInfo | None = None,
     ) -> None:
         """Initialize the integration sensor."""
         self._attr_unique_id = unique_id
@@ -211,6 +239,7 @@ class IntegrationSensor(RestoreSensor):
         self._attr_icon = "mdi:chart-histogram"
         self._source_entity: str = source_entity
         self._last_valid_state: Decimal | None = None
+        self._attr_device_info = device_info
 
     def _unit(self, source_unit: str) -> str:
         """Derive unit from the source sensor, SI prefix and time unit."""
diff --git a/tests/components/integration/test_sensor.py b/tests/components/integration/test_sensor.py
index 355d13c84d6..5b3734bd1be 100644
--- a/tests/components/integration/test_sensor.py
+++ b/tests/components/integration/test_sensor.py
@@ -4,6 +4,7 @@ from datetime import timedelta
 from freezegun import freeze_time
 import pytest
 
+from homeassistant.components.integration.const import DOMAIN
 from homeassistant.components.sensor import SensorDeviceClass, SensorStateClass
 from homeassistant.const import (
     ATTR_UNIT_OF_MEASUREMENT,
@@ -16,10 +17,15 @@ from homeassistant.const import (
     UnitOfTime,
 )
 from homeassistant.core import HomeAssistant, State
+from homeassistant.helpers import device_registry as dr, entity_registry as er
 from homeassistant.setup import async_setup_component
 import homeassistant.util.dt as dt_util
 
-from tests.common import mock_restore_cache, mock_restore_cache_with_extra_data
+from tests.common import (
+    MockConfigEntry,
+    mock_restore_cache,
+    mock_restore_cache_with_extra_data,
+)
 
 
 @pytest.mark.parametrize("method", ["trapezoidal", "left", "right"])
@@ -671,3 +677,47 @@ async def test_calc_errors(hass: HomeAssistant, method) -> None:
     state = hass.states.get("sensor.integration")
     assert state is not None
     assert round(float(state.state)) == 0 if method != "right" else 1
+
+
+async def test_device_id(hass: HomeAssistant) -> None:
+    """Test for source entity device for Riemann sum integral."""
+    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
+
+    integration_config_entry = MockConfigEntry(
+        data={},
+        domain=DOMAIN,
+        options={
+            "method": "trapezoidal",
+            "name": "integration",
+            "round": 1.0,
+            "source": "sensor.test_source",
+            "unit_prefix": "k",
+            "unit_time": "min",
+        },
+        title="Integration",
+    )
+
+    integration_config_entry.add_to_hass(hass)
+
+    assert await hass.config_entries.async_setup(integration_config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    integration_entity = entity_registry.async_get("sensor.integration")
+    assert integration_entity is not None
+    assert integration_entity.device_id == source_entity.device_id