diff --git a/custom_components/aquanta/sensor.py b/custom_components/aquanta/sensor.py
index f6b468c..fbbd86c 100644
--- a/custom_components/aquanta/sensor.py
+++ b/custom_components/aquanta/sensor.py
@@ -8,7 +8,7 @@
     SensorEntityDescription,
 )
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import TEMP_CELSIUS, PERCENTAGE
+from homeassistant.const import UnitOfTemperature, PERCENTAGE
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
@@ -23,7 +23,7 @@
             name="Temperature",
             device_class=SensorDeviceClass.TEMPERATURE,
             state_class=SensorStateClass.MEASUREMENT,
-            native_unit_of_measurement=TEMP_CELSIUS,
+            native_unit_of_measurement=UnitOfTemperature.CELSIUS,
             icon="mdi:water-thermometer",
         ),
         "native_value": lambda entity: entity.coordinator.data["devices"][
@@ -38,7 +38,7 @@
             name="Set point",
             device_class=SensorDeviceClass.TEMPERATURE,
             state_class=SensorStateClass.MEASUREMENT,
-            native_unit_of_measurement=TEMP_CELSIUS,
+            native_unit_of_measurement=UnitOfTemperature.CELSIUS,
             icon="mdi:thermometer-water",
         ),
         "native_value": lambda entity: entity.coordinator.data["devices"][
diff --git a/custom_components/aquanta/water_heater.py b/custom_components/aquanta/water_heater.py
index e9bd814..05ac990 100644
--- a/custom_components/aquanta/water_heater.py
+++ b/custom_components/aquanta/water_heater.py
@@ -9,9 +9,7 @@
     WaterHeaterEntityFeature,
 )
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import (
-    TEMP_CELSIUS,
-)
+from homeassistant.const import UnitOfTemperature
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
@@ -41,7 +39,7 @@ class AquantaWaterHeater(AquantaEntity, WaterHeaterEntity):
 
     _attr_has_entity_name = True
     _attr_supported_features = WaterHeaterEntityFeature.AWAY_MODE
-    _attr_temperature_unit = TEMP_CELSIUS
+    _attr_temperature_unit = UnitOfTemperature.CELSIUS
     _attr_operation_list = [STATE_ECO, STATE_PERFORMANCE, STATE_OFF]
     _attr_name = "Water heater"