From 640eb792fd3b13a9a3181b23cf7cbb7607dd5321 Mon Sep 17 00:00:00 2001
From: fustom <fustom@gmail.com>
Date: Sat, 26 Aug 2023 23:42:22 +0200
Subject: [PATCH] Support whe type 7 (#76)

---
 ariston/__init__.py    |  9 +++++++--
 ariston/ariston_api.py | 14 ++++++++++++++
 ariston/const.py       |  1 +
 ariston/lux2_device.py | 22 ++++++++++++++++++++++
 4 files changed, 44 insertions(+), 2 deletions(-)
 create mode 100644 ariston/lux2_device.py

diff --git a/ariston/__init__.py b/ariston/__init__.py
index 9b42daa..c21dfa9 100644
--- a/ariston/__init__.py
+++ b/ariston/__init__.py
@@ -3,8 +3,6 @@
 import logging
 from typing import Any, Optional
 
-from ariston.bsb_device import AristonBsbDevice
-
 from .ariston_api import AristonAPI, ConnectionException
 from .const import (
     DeviceAttribute,
@@ -12,7 +10,9 @@
     VelisDeviceAttribute,
     WheType,
 )
+from .bsb_device import AristonBsbDevice
 from .lux_device import AristonLuxDevice
+from .lux2_device import AristonLux2Device
 from .evo_device import AristonEvoDevice
 from .galevo_device import AristonGalevoDevice
 from .lydos_hybrid_device import AristonLydosHybridDevice
@@ -110,6 +110,11 @@ def _get_device(
                 api,
                 device,
             )
+        if whe_type == WheType.Lux2.value:
+            return AristonLux2Device(
+                api,
+                device,
+            )
         _LOGGER.exception("Unsupported whe type %s", whe_type)
         return None
 
diff --git a/ariston/ariston_api.py b/ariston/ariston_api.py
index 13d15d1..cb3812e 100644
--- a/ariston/ariston_api.py
+++ b/ariston/ariston_api.py
@@ -347,6 +347,13 @@ def set_evo_eco_mode(self, gw_id: str, eco_mode: bool) -> None:
             eco_mode,
         )
 
+    def set_lux_power_option(self, gw_id: str, power_option: bool) -> None:
+        """Set Velis Lux2 power option"""
+        self._post(
+            f"{ARISTON_API_URL}{ARISTON_VELIS}/{PlantData.Med.value}/{gw_id}/switchPowerOption",
+            power_option,
+        )
+
     def set_velis_power(self, plant_data: PlantData, gw_id: str, power: bool) -> None:
         """Set Velis power"""
         self._post(
@@ -736,6 +743,13 @@ async def async_set_evo_eco_mode(self, gw_id: str, eco_mode: bool) -> None:
             eco_mode,
         )
 
+    async def async_set_lux_power_option(self, gw_id: str, power_option: bool) -> None:
+        """Set Velis Lux2 power option"""
+        await self._async_post(
+            f"{ARISTON_API_URL}{ARISTON_VELIS}/{PlantData.Med.value}/{gw_id}/switchPowerOption",
+            power_option,
+        )
+
     async def async_set_velis_power(
         self, plant_data: PlantData, gw_id: str, power: bool
     ) -> None:
diff --git a/ariston/const.py b/ariston/const.py
index 90dc13e..2432107 100644
--- a/ariston/const.py
+++ b/ariston/const.py
@@ -238,6 +238,7 @@ class WheType(Enum):
     NuosSplit = 4
     Andris2 = 5
     Evo2 = 6
+    Lux2 = 7
     Lux = 8
 
 
diff --git a/ariston/lux2_device.py b/ariston/lux2_device.py
new file mode 100644
index 0000000..b29ca10
--- /dev/null
+++ b/ariston/lux2_device.py
@@ -0,0 +1,22 @@
+"""Lux2 device class for Ariston module."""
+from __future__ import annotations
+
+import logging
+
+from .const import EvoDeviceProperties
+from .evo_device import AristonEvoDevice
+
+_LOGGER = logging.getLogger(__name__)
+
+class AristonLux2Device(AristonEvoDevice):
+    """Class representing a physical device, it's state and properties."""
+
+    def set_water_heater_power_option(self, power_option: bool):
+        """Set water heater power option"""
+        self.api.set_lux_power_option(self.gw, power_option)
+        self.data[EvoDeviceProperties.PWR_OPT] = power_option
+
+    async def async_set_water_heater_power_option(self, power_option: bool):
+        """Async set water heater power option"""
+        await self.api.async_set_lux_power_option(self.gw, power_option)
+        self.data[EvoDeviceProperties.PWR_OPT] = power_option