diff --git a/custom_components/powercalc/sensors/energy.py b/custom_components/powercalc/sensors/energy.py
index c2a830220..c6212fc68 100644
--- a/custom_components/powercalc/sensors/energy.py
+++ b/custom_components/powercalc/sensors/energy.py
@@ -59,57 +59,89 @@ async def create_energy_sensor(
     source_entity: SourceEntity | None = None,
 ) -> EnergySensor:
     """Create the energy sensor entity."""
-    # User specified an existing energy sensor with "energy_sensor_id" option. Just return that one
-    if CONF_ENERGY_SENSOR_ID in sensor_config:
-        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,
+
+    # Check for existing energy sensor
+    energy_sensor = await _get_existing_energy_sensor(hass, sensor_config)
+    if energy_sensor:
+        return energy_sensor
+
+    # Check if we should find or create a related energy sensor
+    energy_sensor = await _get_related_energy_sensor(hass, sensor_config, power_sensor)
+    if energy_sensor:
+        return energy_sensor
+
+    # Create a new virtual energy sensor based on the virtual power sensor
+    return await _create_virtual_energy_sensor(hass, sensor_config, power_sensor, source_entity)
+
+
+async def _get_existing_energy_sensor(
+    hass: HomeAssistant,
+    sensor_config: ConfigType,
+) -> EnergySensor | None:
+    """Check if the user specified an 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,
+    )
 
-    # User specified an existing power sensor with "power_sensor_id" option. Try to find a corresponding energy sensor
-    if CONF_POWER_SENSOR_ID in sensor_config and isinstance(
-        power_sensor,
-        RealPowerSensor,
-    ):
-        # User can force the energy sensor creation with "force_energy_sensor_creation" option.
-        # If they did, don't look for an energy sensor
-        if CONF_FORCE_ENERGY_SENSOR_CREATION not in sensor_config or not sensor_config.get(CONF_FORCE_ENERGY_SENSOR_CREATION):
-            real_energy_sensor = find_related_real_energy_sensor(hass, power_sensor)
-            if real_energy_sensor:
-                _LOGGER.debug(
-                    "Found existing energy sensor '%s' for the power sensor '%s'",
-                    real_energy_sensor.entity_id,
-                    power_sensor.entity_id,
-                )
-                return real_energy_sensor  # type: ignore
-            _LOGGER.debug(
-                "No existing energy sensor found for the power sensor '%s'",
-                power_sensor.entity_id,
-            )
-        else:
-            _LOGGER.debug(
-                "Forced energy sensor generation for the power sensor '%s'",
-                power_sensor.entity_id,
-            )
-
-    # Create an energy sensor based on riemann integral integration, which uses the virtual powercalc sensor as source.
+
+async def _get_related_energy_sensor(
+    hass: HomeAssistant,
+    sensor_config: ConfigType,
+    power_sensor: PowerSensor,
+) -> EnergySensor | None:
+    """Find or create a related energy sensor based on the power sensor."""
+
+    if CONF_POWER_SENSOR_ID not in sensor_config or not isinstance(power_sensor, RealPowerSensor):
+        return None
+
+    if sensor_config.get(CONF_FORCE_ENERGY_SENSOR_CREATION):
+        _LOGGER.debug(
+            "Forced energy sensor generation for the power sensor '%s'",
+            power_sensor.entity_id,
+        )
+        return None
+
+    real_energy_sensor = _find_related_real_energy_sensor(hass, power_sensor)
+    if real_energy_sensor:
+        _LOGGER.debug(
+            "Found existing energy sensor '%s' for the power sensor '%s'",
+            real_energy_sensor.entity_id,
+            power_sensor.entity_id,
+        )
+        return real_energy_sensor  # type: ignore
+
+    _LOGGER.debug(
+        "No existing energy sensor found for the power sensor '%s'",
+        power_sensor.entity_id,
+    )
+    return None
+
+
+async def _create_virtual_energy_sensor(
+    hass: HomeAssistant,
+    sensor_config: ConfigType,
+    power_sensor: PowerSensor,
+    source_entity: SourceEntity | None,
+) -> VirtualEnergySensor:
+    """Create a virtual energy sensor using riemann integral integration."""
     name = generate_energy_sensor_name(
         sensor_config,
         sensor_config.get(CONF_NAME),
         source_entity,
     )
-    unique_id = None
-    if power_sensor.unique_id:
-        unique_id = f"{power_sensor.unique_id}_energy"
-
+    unique_id = f"{power_sensor.unique_id}_energy" if power_sensor.unique_id else None
     entity_id = generate_energy_sensor_entity_id(
         hass,
         sensor_config,
@@ -117,7 +149,6 @@ async def create_energy_sensor(
         unique_id=unique_id,
     )
     entity_category = sensor_config.get(CONF_ENERGY_SENSOR_CATEGORY)
-
     unit_prefix = get_unit_prefix(hass, sensor_config, power_sensor)
 
     _LOGGER.debug(
@@ -164,7 +195,7 @@ def get_unit_prefix(
 
 
 @callback
-def find_related_real_energy_sensor(
+def _find_related_real_energy_sensor(
     hass: HomeAssistant,
     power_sensor: RealPowerSensor,
 ) -> RealEnergySensor | None: