Skip to content

Commit

Permalink
feat: Updated cost tracker entities to be associated with device of t…
Browse files Browse the repository at this point in the history
…racked entity, if one exists (1 hour dev time)
  • Loading branch information
BottlecapDave committed Dec 2, 2024
1 parent 88e25a8 commit 9fe69a0
Show file tree
Hide file tree
Showing 5 changed files with 50 additions and 13 deletions.
15 changes: 15 additions & 0 deletions custom_components/octopus_energy/cost_tracker/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
6 changes: 4 additions & 2 deletions custom_components/octopus_energy/cost_tracker/cost_tracker.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)
Expand All @@ -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.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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

Expand All @@ -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.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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

Expand All @@ -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.
Expand Down
30 changes: 23 additions & 7 deletions custom_components/octopus_energy/sensor.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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)
Expand All @@ -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
Expand Down

0 comments on commit 9fe69a0

Please sign in to comment.