fix: cleanup duplication in energy sensor resolving

Author: bramstroker

custom_components/powercalc/sensors/cost.py

@@ -13,7 +13,6 @@
     UnitOfEnergy,
 )
 from homeassistant.core import Event, EventStateChangedData, HomeAssistant, State, callback
-import homeassistant.helpers.entity_registry as er
 from homeassistant.helpers.event import async_track_state_change_event
 from homeassistant.helpers.restore_state import RestoreEntity
 from homeassistant.helpers.typing import ConfigType
@@ -32,15 +31,14 @@
     DOMAIN,
     DOMAIN_CONFIG,
 )
-from custom_components.powercalc.errors import SensorConfigurationError
 from custom_components.powercalc.unit import convert_to_decimal, parse_decimal
 
 from .abstract import (
     BaseEntity,
     generate_cost_sensor_entity_id,
     generate_cost_sensor_name,
 )
-from .energy import EnergySensor, RealEnergySensor
+from .energy import EnergySensor, resolve_existing_energy_sensor
 
 if TYPE_CHECKING:
     from datetime import datetime
@@ -188,17 +186,7 @@ def create_cost_sensor_for_energy_entity(hass: HomeAssistant, sensor_config: Con
     """
     cost_config = sensor_config.get(CONF_COST, {})
     energy_sensor_id = cost_config.get(CONF_ENERGY_SENSOR_ID) or sensor_config[CONF_ENERGY_SENSOR_ID]
-    entity_entry = er.async_get(hass).async_get(energy_sensor_id)
-    if entity_entry is None:
-        raise SensorConfigurationError(
-            f"No energy sensor with id {energy_sensor_id} found in your HA instance. "
-            "Double check the `energy_sensor_id` setting",
-        )
-    energy_sensor = RealEnergySensor(
-        entity_entry.entity_id,
-        entity_entry.name or entity_entry.original_name,
-        entity_entry.unique_id,
-    )
+    energy_sensor = resolve_existing_energy_sensor(hass, energy_sensor_id)
     # In YAML the name is optional, fall back to the name of the tracked energy sensor.
     name = sensor_config.get(CONF_NAME) or energy_sensor.name
     return create_cost_sensor(

custom_components/powercalc/sensors/energy.py

@@ -80,6 +80,17 @@ def create_energy_sensor(
     return _create_virtual_energy_sensor(hass, sensor_config, power_sensor, source_entity)
 
 
+def resolve_existing_energy_sensor(hass: HomeAssistant, energy_sensor_id: str) -> RealEnergySensor:
+    """Look up an existing energy sensor in the entity registry, raising when not found."""
+    entity_entry = er.async_get(hass).async_get(energy_sensor_id)
+    if entity_entry is None:
+        raise SensorConfigurationError(
+            f"No energy sensor with id {energy_sensor_id} found in your HA instance. "
+            "Double check the `energy_sensor_id` setting",
+        )
+    return RealEnergySensor.from_registry_entry(entity_entry)
+
+
 def _get_existing_energy_sensor(
     hass: HomeAssistant,
     sensor_config: ConfigType,
@@ -88,19 +99,7 @@ def _get_existing_energy_sensor(
     if CONF_ENERGY_SENSOR_ID not in sensor_config:
         return None
 
-    ent_reg = er.async_get(hass)
-    energy_sensor_id = sensor_config[CONF_ENERGY_SENSOR_ID]
-    entity_entry = ent_reg.async_get(energy_sensor_id)
-    if entity_entry is None:
-        raise SensorConfigurationError(
-            f"No energy sensor with id {energy_sensor_id} found in your HA instance. "
-            "Double check `energy_sensor_id` setting",
-        )
-    return RealEnergySensor(
-        entity_entry.entity_id,
-        entity_entry.name or entity_entry.original_name,
-        entity_entry.unique_id,
-    )
+    return resolve_existing_energy_sensor(hass, sensor_config[CONF_ENERGY_SENSOR_ID])
 
 
 def _get_related_energy_sensor(
@@ -227,12 +226,7 @@ def _find_related_real_energy_sensor(
     if not energy_sensors:
         return None
 
-    entity_entry = energy_sensors[0]
-    return RealEnergySensor(
-        entity_entry.entity_id,
-        entity_entry.name or entity_entry.original_name,
-        entity_entry.unique_id,
-    )
+    return RealEnergySensor.from_registry_entry(energy_sensors[0])
 
 
 class EnergySensor(BaseEntity):
@@ -365,6 +359,11 @@ def __init__(
         self._name = name
         self._unique_id = unique_id
 
+    @classmethod
+    def from_registry_entry(cls, entry: er.RegistryEntry) -> RealEnergySensor:
+        """Create a reference to an existing energy sensor from its registry entry."""
+        return cls(entry.entity_id, entry.name or entry.original_name, entry.unique_id)
+
     @property
     def name(self) -> str | None:
         """Return the name of the sensor."""