From 0bc4c443ca0559dd804c99b43fb3c68866b5416e Mon Sep 17 00:00:00 2001 From: dougiteixeira <31328123+dougiteixeira@users.noreply.github.com> Date: Fri, 16 Jun 2023 14:01:16 +0000 Subject: [PATCH 1/3] Adds the device of the original entity in the Integration helper sensor - Riemann sum integral sensor (integration) --- .../components/integration/sensor.py | 31 ++++++++++++++++++- 1 file changed, 30 insertions(+), 1 deletion(-) diff --git a/homeassistant/components/integration/sensor.py b/homeassistant/components/integration/sensor.py index b28b426d3af3..2ef80aa10cd5 100644 --- a/homeassistant/components/integration/sensor.py +++ b/homeassistant/components/integration/sensor.py @@ -28,7 +28,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 @@ -144,6 +149,27 @@ async def async_setup_entry( if unit_prefix == "none": unit_prefix = None + 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 + integral = IntegrationSensor( integration_method=config_entry.options[CONF_METHOD], name=config_entry.title, @@ -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 @@ def __init__( 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 @@ def __init__( 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.""" From 7e6c5a128527715aa98e73cf8d89e3b8f581890e Mon Sep 17 00:00:00 2001 From: dougiteixeira <31328123+dougiteixeira@users.noreply.github.com> Date: Fri, 16 Jun 2023 22:43:36 +0000 Subject: [PATCH 2/3] Add test for sensor.py --- tests/components/integration/test_sensor.py | 52 ++++++++++++++++++++- 1 file changed, 51 insertions(+), 1 deletion(-) diff --git a/tests/components/integration/test_sensor.py b/tests/components/integration/test_sensor.py index 355d13c84d6c..5d4f0d59a845 100644 --- a/tests/components/integration/test_sensor.py +++ b/tests/components/integration/test_sensor.py @@ -4,6 +4,7 @@ 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 @@ 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 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 + + 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 From 66487538c34d72599306343ee771d18dd66c25cb Mon Sep 17 00:00:00 2001 From: dougiteixeira <31328123+dougiteixeira@users.noreply.github.com> Date: Sat, 17 Jun 2023 01:06:14 +0000 Subject: [PATCH 3/3] Update --- homeassistant/components/integration/sensor.py | 10 +++++----- tests/components/integration/test_sensor.py | 2 +- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/homeassistant/components/integration/sensor.py b/homeassistant/components/integration/sensor.py index 2ef80aa10cd5..ad0f96dd540a 100644 --- a/homeassistant/components/integration/sensor.py +++ b/homeassistant/components/integration/sensor.py @@ -145,13 +145,9 @@ async def async_setup_entry( registry, config_entry.options[CONF_SOURCE_SENSOR] ) - unit_prefix = config_entry.options[CONF_UNIT_PREFIX] - if unit_prefix == "none": - unit_prefix = None - 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) @@ -170,6 +166,10 @@ async def async_setup_entry( else: device_info = None + unit_prefix = config_entry.options[CONF_UNIT_PREFIX] + if unit_prefix == "none": + unit_prefix = None + integral = IntegrationSensor( integration_method=config_entry.options[CONF_METHOD], name=config_entry.title, diff --git a/tests/components/integration/test_sensor.py b/tests/components/integration/test_sensor.py index 5d4f0d59a845..5b3734bd1be9 100644 --- a/tests/components/integration/test_sensor.py +++ b/tests/components/integration/test_sensor.py @@ -680,7 +680,7 @@ async def test_calc_errors(hass: HomeAssistant, method) -> None: async def test_device_id(hass: HomeAssistant) -> None: - """Test integration sensor units using a power source.""" + """Test for source entity device for Riemann sum integral.""" device_registry = dr.async_get(hass) entity_registry = er.async_get(hass)