From 76f0b306b8fa2213bb1c5df8ed4a00e10b1da358 Mon Sep 17 00:00:00 2001
From: Guy Badman <61918526+badguy99@users.noreply.github.com>
Date: Thu, 21 May 2020 07:11:13 +0100
Subject: [PATCH] maintain backwards compatibility

---
 custom_components/badnest/water_heater.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/custom_components/badnest/water_heater.py b/custom_components/badnest/water_heater.py
index ca1916a..9d20c06 100644
--- a/custom_components/badnest/water_heater.py
+++ b/custom_components/badnest/water_heater.py
@@ -11,7 +11,6 @@ from homeassistant.const import (
     ATTR_ENTITY_ID,
 )
 from homeassistant.components.water_heater import (
-    WaterHeaterDevice,
     SCAN_INTERVAL,
     STATE_OFF,
     STATE_ON,
@@ -22,6 +21,11 @@ from homeassistant.components.water_heater import (
     ATTR_OPERATION_MODE,
     ATTR_OPERATION_LIST,
 )
+try:
+    from homeassistant.components.water_heater import WaterHeaterEntity
+except ImportError:
+    from homeassistant.components.water_heater import WaterHeaterDevice as WaterHeaterEntity
+
 from .const import (
     DOMAIN,
 )