Get rid off unused entities

4 files changed, 34 insertions, 22 deletions

diff --git a/custom_components/skykettle/light.py b/custom_components/skykettle/light.py
index 11e3f26..cb0628d 100644
--- a/custom_components/skykettle/light.py
+++ b/custom_components/skykettle/light.py
@@ -22,15 +22,17 @@ LIGHT_GAME = "light"
 
 async def async_setup_entry(hass, entry, async_add_entities, discovery_info=None):
     """Set up the SkyKettle entry."""
-    async_add_entities([
-        KettleLight(hass, entry, LIGHT_GAME),
-        KettleLight(hass, entry, SkyKettle.LIGHT_BOIL, 0),
-        KettleLight(hass, entry, SkyKettle.LIGHT_BOIL, 1),
-        KettleLight(hass, entry, SkyKettle.LIGHT_BOIL, 2),
-        KettleLight(hass, entry, SkyKettle.LIGHT_LAMP, 0),
-        KettleLight(hass, entry, SkyKettle.LIGHT_LAMP, 1),
-        KettleLight(hass, entry, SkyKettle.LIGHT_LAMP, 2),
-    ])
+    model_code = hass.data[DOMAIN][entry.entry_id][DATA_CONNECTION].model_code
+    if model_code in [1, 2]:
+        async_add_entities([
+            KettleLight(hass, entry, LIGHT_GAME),
+            KettleLight(hass, entry, SkyKettle.LIGHT_BOIL, 0),
+            KettleLight(hass, entry, SkyKettle.LIGHT_BOIL, 1),
+            KettleLight(hass, entry, SkyKettle.LIGHT_BOIL, 2),
+            KettleLight(hass, entry, SkyKettle.LIGHT_LAMP, 0),
+            KettleLight(hass, entry, SkyKettle.LIGHT_LAMP, 1),
+            KettleLight(hass, entry, SkyKettle.LIGHT_LAMP, 2),
+        ])
 
 
 class KettleLight(LightEntity):
diff --git a/custom_components/skykettle/number.py b/custom_components/skykettle/number.py
index 27d590c..bb9f7f7 100644
--- a/custom_components/skykettle/number.py
+++ b/custom_components/skykettle/number.py
@@ -22,14 +22,16 @@ NUMBER_LAMP_AUTO_OFF_HOURS = "lamp_auto_off_hours"
 
 async def async_setup_entry(hass, entry, async_add_entities):
     """Set up the SkyKettle entry."""
-    async_add_entities([
-        SkyNumber(hass, entry, NUMBER_TYPE_BOIL_TIME),
-        SkyNumber(hass, entry, NUMBER_TEMPERATURE_LOW),
-        SkyNumber(hass, entry, NUMBER_TEMPERATURE_MID),
-        SkyNumber(hass, entry, NUMBER_TEMPERATURE_HIGH),
-        SkyNumber(hass, entry, NUMBER_COLOR_INTERVAL),
-        SkyNumber(hass, entry, NUMBER_LAMP_AUTO_OFF_HOURS),
-    ])
+    model_code = hass.data[DOMAIN][entry.entry_id][DATA_CONNECTION].model_code
+    if model_code in [1, 2]:
+        async_add_entities([
+            SkyNumber(hass, entry, NUMBER_TYPE_BOIL_TIME),
+            SkyNumber(hass, entry, NUMBER_TEMPERATURE_LOW),
+            SkyNumber(hass, entry, NUMBER_TEMPERATURE_MID),
+            SkyNumber(hass, entry, NUMBER_TEMPERATURE_HIGH),
+            SkyNumber(hass, entry, NUMBER_COLOR_INTERVAL),
+            SkyNumber(hass, entry, NUMBER_LAMP_AUTO_OFF_HOURS),
+        ])
 
 
 class SkyNumber(NumberEntity):
diff --git a/custom_components/skykettle/sensor.py b/custom_components/skykettle/sensor.py
index af8089d..3c589d9 100644
--- a/custom_components/skykettle/sensor.py
+++ b/custom_components/skykettle/sensor.py
@@ -18,10 +18,14 @@ SENSOR_TYPE_SUCCESS_RATE = "success_rate"
 
 async def async_setup_entry(hass, entry, async_add_entities):
     """Set up the SkyKettle entry."""
+    model_code = hass.data[DOMAIN][entry.entry_id][DATA_CONNECTION].model_code
     async_add_entities([
-        SkySensor(hass, entry, SENSOR_TYPE_WATER_FRESHNESS),
         SkySensor(hass, entry, SENSOR_TYPE_SUCCESS_RATE),
     ])
+    if model_code in [1, 2]:
+        async_add_entities([
+            SkySensor(hass, entry, SENSOR_TYPE_WATER_FRESHNESS),
+        ])
 
 
 class SkySensor(SensorEntity):
diff --git a/custom_components/skykettle/switch.py b/custom_components/skykettle/switch.py
index 805b0b9..64c73b6 100644
--- a/custom_components/skykettle/switch.py
+++ b/custom_components/skykettle/switch.py
@@ -20,12 +20,16 @@ SWITCH_LIGHT_BOIL = "light_boil"
 
 async def async_setup_entry(hass, entry, async_add_entities):
     """Set up the SkyKettle entry."""
+    model_code = hass.data[DOMAIN][entry.entry_id][DATA_CONNECTION].model_code
     async_add_entities([
-        SkySwitch(hass, entry, SWITCH_MAIN),
-        SkySwitch(hass, entry, SWITCH_SOUND),
-        SkySwitch(hass, entry, SWITCH_LIGHT_SYNC),
-        SkySwitch(hass, entry, SWITCH_LIGHT_BOIL),
+        SkySwitch(hass, entry, SWITCH_MAIN)
     ])
+    if model_code in [1, 2]:
+        async_add_entities([
+            SkySwitch(hass, entry, SWITCH_SOUND),
+            SkySwitch(hass, entry, SWITCH_LIGHT_SYNC),
+            SkySwitch(hass, entry, SWITCH_LIGHT_BOIL),
+        ])
 
 
 class SkySwitch(SwitchEntity):