Add the device of the source entity in the helper entities for Utility Meter (#94734)

Co-authored-by: Franck Nijhof <frenck@frenck.nl>
pull/95286/head
dougiteixeira 2023-06-26 13:08:13 -03:00 committed by GitHub
parent 403496eb92
commit 39229ce098
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 158 additions and 5 deletions

View File

@ -7,11 +7,22 @@ from homeassistant.components.select import SelectEntity
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import CONF_UNIQUE_ID
from homeassistant.core import HomeAssistant
from homeassistant.helpers import (
device_registry as dr,
entity_registry as er,
)
from homeassistant.helpers.entity import DeviceInfo
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.restore_state import RestoreEntity
from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
from .const import CONF_METER, CONF_TARIFFS, DATA_UTILITY, TARIFF_ICON
from .const import (
CONF_METER,
CONF_SOURCE_SENSOR,
CONF_TARIFFS,
DATA_UTILITY,
TARIFF_ICON,
)
_LOGGER = logging.getLogger(__name__)
@ -26,7 +37,35 @@ async def async_setup_entry(
tariffs: list[str] = config_entry.options[CONF_TARIFFS]
unique_id = config_entry.entry_id
tariff_select = TariffSelect(name, tariffs, unique_id)
registry = er.async_get(hass)
source_entity = registry.async_get(config_entry.options[CONF_SOURCE_SENSOR])
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
tariff_select = TariffSelect(
name,
tariffs,
unique_id,
device_info=device_info,
)
async_add_entities([tariff_select])
@ -63,10 +102,17 @@ async def async_setup_platform(
class TariffSelect(SelectEntity, RestoreEntity):
"""Representation of a Tariff selector."""
def __init__(self, name, tariffs, unique_id):
def __init__(
self,
name,
tariffs,
unique_id,
device_info: DeviceInfo | None = None,
) -> None:
"""Initialize a tariff selector."""
self._attr_name = name
self._attr_unique_id = unique_id
self._attr_device_info = device_info
self._current_tariff: str | None = None
self._tariffs = tariffs
self._attr_icon = TARIFF_ICON

View File

@ -28,8 +28,13 @@ from homeassistant.const import (
UnitOfEnergy,
)
from homeassistant.core import Event, HomeAssistant, State, callback
from homeassistant.helpers import entity_platform, entity_registry as er
from homeassistant.helpers import (
device_registry as dr,
entity_platform,
entity_registry as er,
)
from homeassistant.helpers.dispatcher import async_dispatcher_connect
from homeassistant.helpers.entity import DeviceInfo
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.event import (
async_track_point_in_time,
@ -120,6 +125,27 @@ async def async_setup_entry(
registry, config_entry.options[CONF_SOURCE_SENSOR]
)
source_entity = registry.async_get(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
cron_pattern = None
delta_values = config_entry.options[CONF_METER_DELTA_VALUES]
meter_offset = timedelta(days=config_entry.options[CONF_METER_OFFSET])
@ -149,6 +175,7 @@ async def async_setup_entry(
tariff_entity=tariff_entity,
tariff=None,
unique_id=entry_id,
device_info=device_info,
)
meters.append(meter_sensor)
hass.data[DATA_UTILITY][entry_id][DATA_TARIFF_SENSORS].append(meter_sensor)
@ -168,6 +195,7 @@ async def async_setup_entry(
tariff_entity=tariff_entity,
tariff=tariff,
unique_id=f"{entry_id}_{tariff}",
device_info=device_info,
)
meters.append(meter_sensor)
hass.data[DATA_UTILITY][entry_id][DATA_TARIFF_SENSORS].append(meter_sensor)
@ -341,9 +369,11 @@ class UtilityMeterSensor(RestoreSensor):
tariff,
unique_id,
suggested_entity_id=None,
device_info=None,
):
"""Initialize the Utility Meter sensor."""
self._attr_unique_id = unique_id
self._attr_device_info = device_info
self.entity_id = suggested_entity_id
self._parent_meter = parent_meter
self._sensor_source_id = source_entity

View File

@ -41,7 +41,7 @@ from homeassistant.const import (
UnitOfEnergy,
)
from homeassistant.core import CoreState, HomeAssistant, State
from homeassistant.helpers import entity_registry as er
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
@ -1458,3 +1458,80 @@ def test_calculate_adjustment_invalid_new_state(
new_state: State = State(entity_id="sensor.test", state="unknown")
assert mock_sensor.calculate_adjustment(None, new_state) is None
assert "Invalid state unknown" in caplog.text
async def test_device_id(hass: HomeAssistant) -> None:
"""Test for source entity device for Utility Meter."""
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={
"cycle": "monthly",
"delta_values": False,
"name": "Energy",
"net_consumption": False,
"offset": 0,
"periodically_resetting": True,
"source": "sensor.test_source",
"tariffs": ["peak", "offpeak"],
},
title="Energy",
)
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("sensor.energy_peak")
assert utility_meter_entity is not None
assert utility_meter_entity.device_id == source_entity.device_id
utility_meter_entity = entity_registry.async_get("sensor.energy_offpeak")
assert utility_meter_entity is not None
assert utility_meter_entity.device_id == source_entity.device_id
utility_meter_no_tariffs_config_entry = MockConfigEntry(
data={},
domain=DOMAIN,
options={
"cycle": "monthly",
"delta_values": False,
"name": "Energy",
"net_consumption": False,
"offset": 0,
"periodically_resetting": True,
"source": "sensor.test_source",
"tariffs": [],
},
title="Energy",
)
utility_meter_no_tariffs_config_entry.add_to_hass(hass)
assert await hass.config_entries.async_setup(
utility_meter_no_tariffs_config_entry.entry_id
)
await hass.async_block_till_done()
utility_meter_no_tariffs_entity = entity_registry.async_get("sensor.energy")
assert utility_meter_no_tariffs_entity is not None
assert utility_meter_no_tariffs_entity.device_id == source_entity.device_id