diff --git a/custom_components/badnest/water_heater.py b/custom_components/badnest/water_heater.py
index 1fa0e6c..ca1916a 100644
--- a/custom_components/badnest/water_heater.py
+++ b/custom_components/badnest/water_heater.py
@@ -96,7 +96,7 @@ async def async_setup_platform(hass,
     )
 
 
-class NestWaterHeater(WaterHeaterDevice):
+class NestWaterHeater(WaterHeaterEntity):
     """Representation of a Nest water heater device."""
 
     def __init__(self, device_id, api):