diff --git a/custom_components/aton_storage/sensor.py b/custom_components/aton_storage/sensor.py
index 5a1f467..cce6dc6 100644
--- a/custom_components/aton_storage/sensor.py
+++ b/custom_components/aton_storage/sensor.py
@@ -18,7 +18,7 @@
     SensorStateClass,
 )
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import PERCENTAGE, POWER_WATT
+from homeassistant.const import PERCENTAGE, POWER_WATT, ENERGY_WATT_HOUR
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import DiscoveryInfoType