Backward compatibility

2 files changed, 12 insertions, 7 deletions

diff --git a/custom_components/skykettle/__init__.py b/custom_components/skykettle/__init__.py
index deb38bf..257ad98 100644
--- a/custom_components/skykettle/__init__.py
+++ b/custom_components/skykettle/__init__.py
@@ -29,6 +29,15 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry):
     if DOMAIN not in hass.data: hass.data[DOMAIN] = {}
     if entry.entry_id not in hass.data: hass.data[DOMAIN][entry.entry_id] = {}
 
+    # Backward compatibility
+    if (CONF_FRIENDLY_NAME in entry.data
+        and not SkyKettle.get_model_code(entry.data[CONF_FRIENDLY_NAME]) 
+        and SkyKettle.get_model_code(entry.data[CONF_FRIENDLY_NAME] + 'S')):
+        config = dict(entry.data.items())
+        config[CONF_FRIENDLY_NAME] = config[CONF_FRIENDLY_NAME] + 'S'
+        hass.config_entries.async_update_entry(entry, data=config)
+        _LOGGER.info(f"Fixed invalid model name: {config[CONF_FRIENDLY_NAME][:-1]} -> {config[CONF_FRIENDLY_NAME]}")
+        
     kettle = KettleConnection(
         mac=entry.data[CONF_MAC],
         key=entry.data[CONF_PASSWORD],
diff --git a/custom_components/skykettle/skykettle.py b/custom_components/skykettle/skykettle.py
index 3633ea7..8d239cc 100644
--- a/custom_components/skykettle/skykettle.py
+++ b/custom_components/skykettle/skykettle.py
@@ -112,17 +112,13 @@ class SkyKettle():
     def __init__(self, model):
         _LOGGER.info(f"Kettle model: {model}")
         self.model = model
-        self.model_code = self.get_model_code(model, backward_compatibility=True)
+        self.model_code = self.get_model_code(model)
         if not self.model_code:
             raise SkyKettleError("Unknown kettle model")
 
     @staticmethod
-    def get_model_code(model, backward_compatibility=False):
-        if model in SkyKettle.MODEL_TYPE:
-            return SkyKettle.MODEL_TYPE[model]
-        if backward_compatibility and model + 'S' in SkyKettle.MODEL_TYPE:
-            return SkyKettle.MODEL_TYPE[model + 'S']
-        return None
+    def get_model_code(model):
+        return SkyKettle.MODEL_TYPE.get(model, None)
 
     @abstractmethod
     async def command(self, command, params=[]):