From 7065b26c37bf8f4ea68f09e4d1c7aa01e6d3c750 Mon Sep 17 00:00:00 2001
From: dougiteixeira <31328123+dougiteixeira@users.noreply.github.com>
Date: Fri, 16 Jun 2023 23:19:34 +0000
Subject: [PATCH 1/3] Adds the device of the original entity in the helper
 entities for Derivative

---
 homeassistant/components/derivative/sensor.py | 31 +++++++++++-
 tests/components/derivative/test_sensor.py    | 48 +++++++++++++++++++
 2 files changed, 78 insertions(+), 1 deletion(-)

diff --git a/homeassistant/components/derivative/sensor.py b/homeassistant/components/derivative/sensor.py
index e1cc278137cfcd..55bfc5e4be7300 100644
--- a/homeassistant/components/derivative/sensor.py
+++ b/homeassistant/components/derivative/sensor.py
@@ -19,7 +19,12 @@
     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
@@ -86,6 +91,27 @@ async def async_setup_entry(
         registry, config_entry.options[CONF_SOURCE]
     )
 
+    source_entity = er.EntityRegistry.async_get(registry, source_entity_id)
+    dev_reg = dr.async_get(hass)
+    # Validate + resolve devicers id of entity_id
+    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
@@ -99,6 +125,7 @@ async def async_setup_entry(
         unit_of_measurement=None,
         unit_prefix=unit_prefix,
         unit_time=config_entry.options[CONF_UNIT_TIME],
+        device_info=device_info,
     )
 
     async_add_entities([derivative_sensor])
@@ -142,9 +169,11 @@ def __init__(
         unit_prefix: str | None,
         unit_time: UnitOfTime,
         unique_id: str | None,
+        device_info: DeviceInfo | None = None,
     ) -> None:
         """Initialize the derivative sensor."""
         self._attr_unique_id = unique_id
+        self._attr_device_info = device_info
         self._sensor_source_id = source_entity
         self._round_digits = round_digits
         self._state: float | int | Decimal = 0
diff --git a/tests/components/derivative/test_sensor.py b/tests/components/derivative/test_sensor.py
index 8260e5a0ada132..f5fca46af7d0f8 100644
--- a/tests/components/derivative/test_sensor.py
+++ b/tests/components/derivative/test_sensor.py
@@ -5,11 +5,15 @@
 
 from freezegun import freeze_time
 
+from homeassistant.components.derivative.const import DOMAIN
 from homeassistant.const import UnitOfPower, UnitOfTime
 from homeassistant.core import HomeAssistant
+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 MockConfigEntry
+
 
 async def test_state(hass: HomeAssistant) -> None:
     """Test derivative sensor state."""
@@ -342,3 +346,47 @@ async def test_suffix(hass: HomeAssistant) -> None:
 
     # Testing a network speed sensor at 1000 bytes/s over 10s  = 10kbytes/s2
     assert round(float(state.state), config["sensor"]["round"]) == 0.0
+
+
+async def test_device_id(hass: HomeAssistant) -> None:
+    """Test integration sensor units using a power source."""
+    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
+
+    derivative_config_entry = MockConfigEntry(
+        data={},
+        domain=DOMAIN,
+        options={
+            "name": "Derivative",
+            "round": 1.0,
+            "source": "sensor.test_source",
+            "time_window": {"seconds": 0.0},
+            "unit_prefix": "k",
+            "unit_time": "min",
+        },
+        title="Derivative",
+    )
+
+    derivative_config_entry.add_to_hass(hass)
+
+    assert await hass.config_entries.async_setup(derivative_config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    derivative_entity = entity_registry.async_get("sensor.derivative")
+    assert derivative_entity is not None
+    assert derivative_entity.device_id == source_entity.device_id

From 787a907550bc1e3db6c38b59809367647cc3a834 Mon Sep 17 00:00:00 2001
From: dougiteixeira <31328123+dougiteixeira@users.noreply.github.com>
Date: Sat, 17 Jun 2023 01:02:50 +0000
Subject: [PATCH 2/3] Update

---
 homeassistant/components/derivative/sensor.py | 2 +-
 tests/components/derivative/test_sensor.py    | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/homeassistant/components/derivative/sensor.py b/homeassistant/components/derivative/sensor.py
index 55bfc5e4be7300..8141c90158c2f7 100644
--- a/homeassistant/components/derivative/sensor.py
+++ b/homeassistant/components/derivative/sensor.py
@@ -93,7 +93,7 @@ async def async_setup_entry(
 
     source_entity = er.EntityRegistry.async_get(registry, source_entity_id)
     dev_reg = dr.async_get(hass)
-    # Validate + resolve devicers id of entity_id
+    # Resolve source entity device
     if (
         (source_entity is not None)
         and (source_entity.device_id is not None)
diff --git a/tests/components/derivative/test_sensor.py b/tests/components/derivative/test_sensor.py
index f5fca46af7d0f8..513e9597572b30 100644
--- a/tests/components/derivative/test_sensor.py
+++ b/tests/components/derivative/test_sensor.py
@@ -349,7 +349,7 @@ async def test_suffix(hass: HomeAssistant) -> None:
 
 
 async def test_device_id(hass: HomeAssistant) -> None:
-    """Test integration sensor units using a power source."""
+    """Test for source entity device for Derivative."""
     device_registry = dr.async_get(hass)
     entity_registry = er.async_get(hass)
 

From 1650c490424fa81a436e7a1c85f8dd0753f45c31 Mon Sep 17 00:00:00 2001
From: dougiteixeira <31328123+dougiteixeira@users.noreply.github.com>
Date: Fri, 23 Jun 2023 02:51:08 +0000
Subject: [PATCH 3/3] Update

---
 homeassistant/components/derivative/sensor.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/homeassistant/components/derivative/sensor.py b/homeassistant/components/derivative/sensor.py
index 8141c90158c2f7..793e8edc7699c0 100644
--- a/homeassistant/components/derivative/sensor.py
+++ b/homeassistant/components/derivative/sensor.py
@@ -91,7 +91,7 @@ async def async_setup_entry(
         registry, config_entry.options[CONF_SOURCE]
     )
 
-    source_entity = er.EntityRegistry.async_get(registry, source_entity_id)
+    source_entity = registry.async_get(source_entity_id)
     dev_reg = dr.async_get(hass)
     # Resolve source entity device
     if (