From 7e054bc1424ac2a9ba7b5fad78e1598a28ff87f4 Mon Sep 17 00:00:00 2001
From: Vincent Wolsink <ik@vincentwolsink.nl>
Date: Sun, 13 Aug 2023 20:44:10 +0200
Subject: [PATCH] Update water boiler detection. Simplify code.

---
 custom_components/aguaiot/climate.py | 41 ++++++++++++----------------
 custom_components/aguaiot/const.py   |  3 ++
 2 files changed, 20 insertions(+), 24 deletions(-)

diff --git a/custom_components/aguaiot/climate.py b/custom_components/aguaiot/climate.py
index 2e26c7d..aa8545e 100644
--- a/custom_components/aguaiot/climate.py
+++ b/custom_components/aguaiot/climate.py
@@ -32,6 +32,8 @@
     AGUA_STATUS_OFF,
     AGUA_STATUS_ON,
     CURRENT_HVAC_MAP_AGUA_HEAT,
+    DEVICE_TYPE_AIR,
+    DEVICE_TYPE_WATER,
 )
 from py_agua_iot import Error as AguaIOTError
 
@@ -56,6 +58,11 @@ def __init__(self, coordinator, device):
         CoordinatorEntity.__init__(self, coordinator)
         self._device = device
 
+        if getattr(self._device, "water_temperature", 0) > 0:
+            self._device_type = DEVICE_TYPE_WATER
+        else:
+            self._device_type = DEVICE_TYPE_AIR
+
     @property
     def supported_features(self):
         """Return the list of supported features."""
@@ -108,34 +115,22 @@ def temperature_unit(self):
     @property
     def min_temp(self):
         """Return the minimum temperature to set."""
-        min_temp = self._device.min_air_temp
-        if min_temp is None:
-            min_temp = self._device.min_water_temp
-        return min_temp
+        return getattr(self._device, f"min_{self._device_type}_temp")
 
     @property
     def max_temp(self):
         """Return the maximum temperature to set."""
-        max_temp = self._device.max_air_temp
-        if max_temp is None:
-            max_temp = self._device.max_water_temp
-        return max_temp
+        return getattr(self._device, f"max_{self._device_type}_temp")
 
     @property
     def current_temperature(self):
         """Return the current temperature."""
-        temp = self._device.air_temperature
-        if temp is None:
-            temp = self._device.water_temperature
-        return temp
+        return getattr(self._device, f"{self._device_type}_temperature")
 
     @property
     def target_temperature(self):
         """Return the temperature we try to reach."""
-        set_temp = self._device.set_air_temperature
-        if set_temp is None:
-            set_temp = self._device.set_water_temperature
-        return set_temp
+        return getattr(self._device, f"set_{self._device_type}_temperature")
 
     @property
     def hvac_mode(self):
@@ -192,14 +187,12 @@ async def async_set_temperature(self, **kwargs):
             return
 
         try:
-            if self._device.air_temperature is not None:
-                await self.hass.async_add_executor_job(
-                    setattr, self._device, "set_air_temperature", temperature
-                )
-            elif self._device.water_temperature is not None:
-                await self.hass.async_add_executor_job(
-                    setattr, self._device, "set_water_temperature", temperature
-                )
+            await self.hass.async_add_executor_job(
+                setattr,
+                self._device,
+                f"set_{self._device_type}_temperature",
+                temperature,
+            )
             await self.coordinator.async_request_refresh()
         except (ValueError, AguaIOTError) as err:
             _LOGGER.error("Failed to set temperature, error: %s", err)
diff --git a/custom_components/aguaiot/const.py b/custom_components/aguaiot/const.py
index bcb0989..710e371 100644
--- a/custom_components/aguaiot/const.py
+++ b/custom_components/aguaiot/const.py
@@ -38,6 +38,9 @@
     AGUA_STATUS_OFF: CURRENT_HVAC_OFF,
 }
 
+DEVICE_TYPE_AIR = "air"
+DEVICE_TYPE_WATER = "water"
+
 PLATFORMS = [
     Platform.SENSOR,
     Platform.CLIMATE,