new fix for energy sensor created correctly for kW power sensor

custom_components/powercalc/sensors/energy.py

@@ -10,7 +10,7 @@
 from homeassistant.components.integration.sensor import IntegrationSensor
 from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN
 from homeassistant.components.sensor import SensorDeviceClass, SensorStateClass
-from homeassistant.const import CONF_NAME, ENERGY_KILO_WATT_HOUR, TIME_HOURS, UnitOfTime
+from homeassistant.const import ATTR_UNIT_OF_MEASUREMENT, CONF_NAME, ENERGY_KILO_WATT_HOUR, TIME_HOURS, UnitOfPower, UnitOfTime
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity import EntityCategory
 from homeassistant.helpers.typing import ConfigType
@@ -98,9 +98,7 @@ async def create_energy_sensor(
     )
     entity_category = sensor_config.get(CONF_ENERGY_SENSOR_CATEGORY)
 
-    unit_prefix = sensor_config.get(CONF_ENERGY_SENSOR_UNIT_PREFIX)
-    if unit_prefix == UnitPrefix.NONE:
-        unit_prefix = None
+    unit_prefix = get_unit_prefix(hass, sensor_config, power_sensor)
 
     _LOGGER.debug("Creating energy sensor: %s", name)
     return VirtualEnergySensor(
@@ -120,6 +118,24 @@ async def create_energy_sensor(
     )
 
 
+def get_unit_prefix(hass: HomeAssistant, sensor_config: ConfigType, power_sensor: PowerSensor) -> str | None:
+    unit_prefix = sensor_config.get(CONF_ENERGY_SENSOR_UNIT_PREFIX)
+
+    power_unit = power_sensor.unit_of_measurement
+    power_state = hass.states.get(power_sensor.entity_id)
+    if power_unit is None and power_state:
+        power_unit = power_state.attributes.get(ATTR_UNIT_OF_MEASUREMENT)
+
+    # When the power sensor is in kW, we don't want to add an extra k prefix.
+    # As this would result in an energy sensor having kkWh unit, which is obviously invalid
+    if power_unit == UnitOfPower.KILO_WATT and unit_prefix == UnitPrefix.KILO:
+        unit_prefix = UnitPrefix.NONE
+
+    if unit_prefix == UnitPrefix.NONE:
+        unit_prefix = None
+    return unit_prefix
+
+
 @callback
 def find_related_real_energy_sensor(
     hass: HomeAssistant,

tests/sensors/test_energy.py

@@ -1,13 +1,21 @@
 import logging
 from datetime import timedelta
+from unittest.mock import patch
 
 import pytest
+from homeassistant.components.sensor import ATTR_STATE_CLASS, SensorStateClass
 from homeassistant.components.utility_meter.sensor import SensorDeviceClass
 from homeassistant.const import (
+    ATTR_DEVICE_CLASS,
     ATTR_ENTITY_ID,
+    ATTR_UNIT_OF_MEASUREMENT,
     CONF_ENTITIES,
     CONF_ENTITY_ID,
+    CONF_NAME,
+    CONF_UNIQUE_ID,
     EntityCategory,
+    UnitOfEnergy,
+    UnitOfPower,
     UnitOfTime,
 )
 from homeassistant.core import HomeAssistant
@@ -211,7 +219,29 @@ async def test_unit_prefix_none(hass: HomeAssistant) -> None:
     await hass.async_block_till_done()
 
     state_attributes = hass.states.get("sensor.test_energy").attributes
-    assert state_attributes.get("unit_of_measurement") == "Wh"
+    assert state_attributes.get("unit_of_measurement") == UnitOfEnergy.WATT_HOUR
+
+
+async def test_unit_prefix_kwh_default(hass: HomeAssistant) -> None:
+    """By default, unit prefix should be k, resulting in kWh energy sensor created for a W power sensor"""
+    await create_input_boolean(hass)
+
+    await run_powercalc_setup(
+        hass,
+        get_simple_fixed_config("input_boolean.test"),
+    )
+
+    async_fire_time_changed(
+        hass,
+        dt.utcnow() + timedelta(hours=1),
+    )
+
+    hass.states.async_set("sensor.test_power", "50.00")
+
+    await hass.async_block_till_done()
+
+    state_attributes = hass.states.get("sensor.test_energy").attributes
+    assert state_attributes.get("unit_of_measurement") == UnitOfEnergy.KILO_WATT_HOUR
 
 
 async def test_set_entity_category(hass: HomeAssistant) -> None:
@@ -253,3 +283,50 @@ async def test_calibrate_service(hass: HomeAssistant) -> None:
     await hass.async_block_till_done()
 
     assert hass.states.get(entity_id).state == "100.0000"
+
+
+async def test_real_power_sensor_kw(hass: HomeAssistant) -> None:
+    """
+    Test that the riemann integral sensor is correclty created and updated for a kW power sensor
+    Fixes https://github.com/bramstroker/homeassistant-powercalc/issues/1676
+    """
+
+    hass.states.async_set(
+        "sensor.test_power",
+        "100",
+        {
+            ATTR_UNIT_OF_MEASUREMENT: UnitOfPower.KILO_WATT,
+            ATTR_DEVICE_CLASS: SensorDeviceClass.POWER,
+            ATTR_STATE_CLASS: SensorStateClass.MEASUREMENT,
+        },
+    )
+    await hass.async_block_till_done()
+
+    await run_powercalc_setup(
+        hass,
+        {
+            CONF_NAME: "Test",
+            CONF_UNIQUE_ID: "1234353",
+            CONF_POWER_SENSOR_ID: "sensor.test_power",
+        },
+    )
+
+    state = hass.states.get("sensor.test_energy")
+    assert state
+
+    now = dt.utcnow() + timedelta(minutes=60)
+    with patch("homeassistant.util.dt.utcnow", return_value=now):
+        hass.states.async_set(
+            "sensor.test_power",
+            "200",
+            {
+                ATTR_UNIT_OF_MEASUREMENT: UnitOfPower.KILO_WATT,
+                ATTR_DEVICE_CLASS: SensorDeviceClass.POWER,
+                ATTR_STATE_CLASS: SensorStateClass.MEASUREMENT,
+            },
+        )
+        await hass.async_block_till_done()
+
+    state = hass.states.get("sensor.test_energy")
+    assert state
+    assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == UnitOfEnergy.KILO_WATT_HOUR

tests/sensors/test_power.py

@@ -1,9 +1,7 @@
 import logging
 from datetime import timedelta
-from unittest.mock import patch
 
 import pytest
-from homeassistant.components.sensor import ATTR_STATE_CLASS, SensorStateClass
 from homeassistant.components.utility_meter.sensor import SensorDeviceClass
 from homeassistant.components.vacuum import (
     ATTR_BATTERY_LEVEL,
@@ -12,8 +10,6 @@
     STATE_RETURNING,
 )
 from homeassistant.const import (
-    ATTR_DEVICE_CLASS,
-    ATTR_UNIT_OF_MEASUREMENT,
     CONF_ATTRIBUTE,
     CONF_ENTITIES,
     CONF_ENTITY_ID,
@@ -23,7 +19,6 @@
     STATE_ON,
     STATE_UNAVAILABLE,
     STATE_UNKNOWN,
-    UnitOfPower,
 )
 from homeassistant.core import EVENT_HOMEASSISTANT_START, CoreState, HomeAssistant
 from homeassistant.setup import async_setup_component
@@ -455,51 +450,3 @@ async def test_manually_configured_sensor_overrides_profile(
     await hass.async_block_till_done()
 
     assert_entity_state(hass, "sensor.test_123_power", "0.00")
-
-
-async def test_real_power_sensor_kw(hass: HomeAssistant) -> None:
-    """
-    Test that the riemann integral sensor is correclty created and updated for a kW power sensor
-    Fixes https://github.com/bramstroker/homeassistant-powercalc/issues/1676
-    """
-    await run_powercalc_setup(
-        hass,
-        {
-            CONF_NAME: "Test",
-            CONF_UNIQUE_ID: "1234353",
-            CONF_POWER_SENSOR_ID: "sensor.test_power",
-        },
-    )
-
-    hass.states.async_set(
-        "sensor.test_power",
-        "100",
-        {
-            ATTR_UNIT_OF_MEASUREMENT: UnitOfPower.KILO_WATT,
-            ATTR_DEVICE_CLASS: SensorDeviceClass.POWER,
-            ATTR_STATE_CLASS: SensorStateClass.MEASUREMENT,
-        },
-    )
-    await hass.async_block_till_done()
-
-    state = hass.states.get("sensor.test_energy")
-    assert state
-
-    now = dt.utcnow() + timedelta(minutes=60)
-    with patch("homeassistant.util.dt.utcnow", return_value=now):
-        hass.states.async_set(
-            "sensor.test_power",
-            "200",
-            {
-                ATTR_UNIT_OF_MEASUREMENT: UnitOfPower.KILO_WATT,
-                ATTR_DEVICE_CLASS: SensorDeviceClass.POWER,
-                ATTR_STATE_CLASS: SensorStateClass.MEASUREMENT,
-            },
-        )
-        await hass.async_block_till_done()
-
-    state = hass.states.get("sensor.test_energy")
-    assert state
-    # Re-enable assertion after issue is fixed
-    # assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == UnitOfEnergy.KILO_WATT_HOUR
-    # assert state.state == "100.0000"