From 9fe69a081b08f1bec3ecfb4ea3b4ed552cc665ef Mon Sep 17 00:00:00 2001 From: BottlecapDave Date: Mon, 2 Dec 2024 19:02:07 +0000 Subject: [PATCH] feat: Updated cost tracker entities to be associated with device of tracked entity, if one exists (1 hour dev time) --- .../octopus_energy/cost_tracker/__init__.py | 15 ++++++++++ .../cost_tracker/cost_tracker.py | 6 ++-- .../cost_tracker/cost_tracker_month.py | 6 ++-- .../cost_tracker/cost_tracker_week.py | 6 ++-- custom_components/octopus_energy/sensor.py | 30 ++++++++++++++----- 5 files changed, 50 insertions(+), 13 deletions(-) diff --git a/custom_components/octopus_energy/cost_tracker/__init__.py b/custom_components/octopus_energy/cost_tracker/__init__.py index ef0cf52e..c0d81f91 100644 --- a/custom_components/octopus_energy/cost_tracker/__init__.py +++ b/custom_components/octopus_energy/cost_tracker/__init__.py @@ -4,6 +4,21 @@ SensorStateClass, ) +from homeassistant.helpers.entity import DeviceInfo + +def get_device_info_from_device_entry(device_entry): + if device_entry is None: + return None + + return DeviceInfo( + identifiers=device_entry.identifiers, + name=device_entry.name, + connections=device_entry.connections, + manufacturer=device_entry.manufacturer, + model=device_entry.model, + sw_version=device_entry.sw_version + ) + class CostTrackerResult: tracked_consumption_data: list untracked_consumption_data: list diff --git a/custom_components/octopus_energy/cost_tracker/cost_tracker.py b/custom_components/octopus_energy/cost_tracker/cost_tracker.py index af357925..f934ce5e 100644 --- a/custom_components/octopus_energy/cost_tracker/cost_tracker.py +++ b/custom_components/octopus_energy/cost_tracker/cost_tracker.py @@ -33,7 +33,7 @@ ) from ..coordinators.electricity_rates import ElectricityRatesCoordinatorResult -from . import add_consumption +from . import add_consumption, get_device_info_from_device_entry from ..electricity import calculate_electricity_consumption_and_cost from ..utils.rate_information import get_rate_index, get_unique_rates from ..utils.attributes import dict_to_typed_dict @@ -43,7 +43,7 @@ class OctopusEnergyCostTrackerSensor(CoordinatorEntity, RestoreSensor): """Sensor for calculating the cost for a given sensor.""" - def __init__(self, hass: HomeAssistant, coordinator, config, peak_type = None): + def __init__(self, hass: HomeAssistant, coordinator, config, device_entry, peak_type = None): """Init sensor.""" # Pass coordinator to base class CoordinatorEntity.__init__(self, coordinator) @@ -60,6 +60,8 @@ def __init__(self, hass: HomeAssistant, coordinator, config, peak_type = None): self._hass = hass self.entity_id = generate_entity_id("sensor.{}", self.unique_id, hass=hass) + self._attr_device_info = get_device_info_from_device_entry(device_entry) + @property def entity_registry_enabled_default(self) -> bool: """Return if the entity should be enabled when first added. diff --git a/custom_components/octopus_energy/cost_tracker/cost_tracker_month.py b/custom_components/octopus_energy/cost_tracker/cost_tracker_month.py index 3429e86b..1fc81eb3 100644 --- a/custom_components/octopus_energy/cost_tracker/cost_tracker_month.py +++ b/custom_components/octopus_energy/cost_tracker/cost_tracker_month.py @@ -29,7 +29,7 @@ DOMAIN, ) -from . import accumulate_cost +from . import accumulate_cost, get_device_info_from_device_entry from ..utils.attributes import dict_to_typed_dict @@ -38,7 +38,7 @@ class OctopusEnergyCostTrackerMonthSensor(RestoreSensor): """Sensor for calculating the cost for a given sensor over the course of a month.""" - def __init__(self, hass: HomeAssistant, config_entry, config, tracked_entity_id: str, peak_type = None): + def __init__(self, hass: HomeAssistant, config_entry, config, device_entry, tracked_entity_id: str, peak_type = None): """Init sensor.""" # Pass coordinator to base class @@ -55,6 +55,8 @@ def __init__(self, hass: HomeAssistant, config_entry, config, tracked_entity_id: self._hass = hass self.entity_id = generate_entity_id("sensor.{}", self.unique_id, hass=hass) + self._attr_device_info = get_device_info_from_device_entry(device_entry) + @property def entity_registry_enabled_default(self) -> bool: """Return if the entity should be enabled when first added. diff --git a/custom_components/octopus_energy/cost_tracker/cost_tracker_week.py b/custom_components/octopus_energy/cost_tracker/cost_tracker_week.py index e77427aa..6e94027d 100644 --- a/custom_components/octopus_energy/cost_tracker/cost_tracker_week.py +++ b/custom_components/octopus_energy/cost_tracker/cost_tracker_week.py @@ -29,7 +29,7 @@ DOMAIN, ) -from . import accumulate_cost +from . import accumulate_cost, get_device_info_from_device_entry from ..utils.attributes import dict_to_typed_dict @@ -38,7 +38,7 @@ class OctopusEnergyCostTrackerWeekSensor(RestoreSensor): """Sensor for calculating the cost for a given sensor over the course of a week.""" - def __init__(self, hass: HomeAssistant, config_entry, config, tracked_entity_id: str, peak_type = None): + def __init__(self, hass: HomeAssistant, config_entry, config, device_entry, tracked_entity_id: str, peak_type = None): """Init sensor.""" # Pass coordinator to base class @@ -55,6 +55,8 @@ def __init__(self, hass: HomeAssistant, config_entry, config, tracked_entity_id: self._hass = hass self.entity_id = generate_entity_id("sensor.{}", self.unique_id, hass=hass) + self._attr_device_info = get_device_info_from_device_entry(device_entry) + @property def entity_registry_enabled_default(self) -> bool: """Return if the entity should be enabled when first added. diff --git a/custom_components/octopus_energy/sensor.py b/custom_components/octopus_energy/sensor.py index 58915b29..0573e8fd 100644 --- a/custom_components/octopus_energy/sensor.py +++ b/custom_components/octopus_energy/sensor.py @@ -5,7 +5,7 @@ from homeassistant.util.dt import (utcnow, now) from homeassistant.core import HomeAssistant -from homeassistant.helpers import entity_platform, issue_registry as ir, entity_registry as er +from homeassistant.helpers import entity_platform, issue_registry as ir, entity_registry as er, device_registry as dr import homeassistant.helpers.config_validation as cv from .electricity.current_consumption import OctopusEnergyCurrentElectricityConsumption @@ -79,6 +79,7 @@ from .const import ( CONFIG_COST_TRACKER_MPAN, CONFIG_ACCOUNT_ID, + CONFIG_COST_TRACKER_TARGET_ENTITY_ID, CONFIG_DEFAULT_LIVE_ELECTRICITY_CONSUMPTION_REFRESH_IN_MINUTES, CONFIG_DEFAULT_LIVE_GAS_CONSUMPTION_REFRESH_IN_MINUTES, CONFIG_KIND, @@ -527,13 +528,28 @@ async def async_setup_cost_sensors(hass: HomeAssistant, entry, config, async_add serial_number = meter["serial_number"] coordinator = hass.data[DOMAIN][account_id][DATA_ELECTRICITY_RATES_COORDINATOR_KEY.format(mpan, serial_number)] - sensor = OctopusEnergyCostTrackerSensor(hass, coordinator, config) + device_registry = dr.async_get(hass) + entity_registry = er.async_get(hass) + + source_entity_id = config[CONFIG_COST_TRACKER_TARGET_ENTITY_ID] + + device_id = None + if source_entity_id: + entity = entity_registry.async_get(source_entity_id) + if entity: + device_id = entity.device_id + + device_entry = None + if device_id is not None: + device_entry = device_registry.async_get(device_id) + + sensor = OctopusEnergyCostTrackerSensor(hass, coordinator, config, device_entry) sensor_entity_id = registry.async_get_entity_id("sensor", DOMAIN, sensor.unique_id) entities = [ sensor, - OctopusEnergyCostTrackerWeekSensor(hass, entry, config, sensor_entity_id if sensor_entity_id is not None else sensor.entity_id), - OctopusEnergyCostTrackerMonthSensor(hass, entry, config, sensor_entity_id if sensor_entity_id is not None else sensor.entity_id), + OctopusEnergyCostTrackerWeekSensor(hass, entry, config, device_entry, sensor_entity_id if sensor_entity_id is not None else sensor.entity_id), + OctopusEnergyCostTrackerMonthSensor(hass, entry, config, device_entry, sensor_entity_id if sensor_entity_id is not None else sensor.entity_id), ] debug_override = await async_get_debug_override(hass, mpan, serial_number) @@ -542,12 +558,12 @@ async def async_setup_cost_sensors(hass: HomeAssistant, entry, config, async_add for unique_rate_index in range(0, total_unique_rates): peak_type = get_peak_type(total_unique_rates, unique_rate_index) if peak_type is not None: - peak_sensor = OctopusEnergyCostTrackerSensor(hass, coordinator, config, peak_type) + peak_sensor = OctopusEnergyCostTrackerSensor(hass, coordinator, config, device_entry, peak_type) peak_sensor_entity_id = registry.async_get_entity_id("sensor", DOMAIN, peak_sensor.unique_id) entities.append(peak_sensor) - entities.append(OctopusEnergyCostTrackerWeekSensor(hass, entry, config, peak_sensor_entity_id if peak_sensor_entity_id is not None else f"sensor.{peak_sensor.unique_id}", peak_type)) - entities.append(OctopusEnergyCostTrackerMonthSensor(hass, entry, config, peak_sensor_entity_id if peak_sensor_entity_id is not None else f"sensor.{peak_sensor.unique_id}", peak_type)) + entities.append(OctopusEnergyCostTrackerWeekSensor(hass, entry, config, device_entry, peak_sensor_entity_id if peak_sensor_entity_id is not None else f"sensor.{peak_sensor.unique_id}", peak_type)) + entities.append(OctopusEnergyCostTrackerMonthSensor(hass, entry, config, device_entry, peak_sensor_entity_id if peak_sensor_entity_id is not None else f"sensor.{peak_sensor.unique_id}", peak_type)) async_add_entities(entities) break