💥 refactor code and test cases

custom_components/vaillant_plus/__init__.py

@@ -10,15 +10,14 @@
 from homeassistant.helpers.typing import ConfigType
 from vaillant_plus_cn_api import Token
 
-from .client import VaillantApiHub, VaillantDeviceApiClient
+from .client import VaillantClient
 from .const import (
-    API_HUB,
+    API_CLIENT,
     CONF_DID,
     CONF_TOKEN,
     DISPATCHERS,
     DOMAIN,
     EVT_TOKEN_UPDATED,
-    WEBSOCKET_CLIENT,
 )
 
 # TODO List the platforms that you want to support.
@@ -36,28 +35,28 @@
 async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     hass.data.setdefault(
         DOMAIN,
-        {API_HUB: VaillantApiHub(hass), DISPATCHERS: {}, WEBSOCKET_CLIENT: {}},
+        {API_CLIENT: {}, DISPATCHERS: {}},
     )
     return True
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up Vaillant Plus from a config entry."""
 
-    hub: VaillantApiHub = hass.data[DOMAIN][API_HUB]
     token = Token.deserialize(entry.data.get(CONF_TOKEN))
     device_id = entry.data.get(CONF_DID)
+    client = VaillantClient(hass, token, device_id)
+
+    async def close_client(_):
+        await client.close()
+
+    hass.data[DOMAIN][API_CLIENT][entry.entry_id] = client
 
     @callback
     def on_token_update(token_new: Token) -> None:
         hass.config_entries.async_update_entry(
             entry, data={**entry.data, CONF_TOKEN: token_new.serialize()}
         )
-        # update hub token
-        hub.update_token(token_new)
-        # should update VaillantDeviceApiClient
-        if entry.entry_id in hass.data[DOMAIN][WEBSOCKET_CLIENT]:
-            hass.loop.create_task(hass.data[DOMAIN][WEBSOCKET_CLIENT][entry.entry_id].update_token(token_new))
 
     unsub = async_dispatcher_connect(
         hass,
@@ -68,20 +67,12 @@ def on_token_update(token_new: Token) -> None:
     hass.data[DOMAIN][DISPATCHERS].setdefault(device_id, [])
     hass.data[DOMAIN][DISPATCHERS][device_id].append(unsub)
 
-    device = await hub.get_device(token, device_id)
-    client = VaillantDeviceApiClient(hass, hub, token, device)
-
-    hass.data[DOMAIN][WEBSOCKET_CLIENT][entry.entry_id] = client
-
-    def close_client(_):
-        return client.close()
-
     unsub_stop = hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, close_client)
     hass.data[DOMAIN][DISPATCHERS][device_id].append(unsub_stop)
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
-    hass.loop.create_task(client.connect())
+    hass.loop.create_task(client.start())
 
     return True
 
@@ -90,14 +81,14 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload a config entry."""
     if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
         if (
-            entry.entry_id in hass.data[DOMAIN][WEBSOCKET_CLIENT]
-            and hass.data[DOMAIN][WEBSOCKET_CLIENT][entry.entry_id] is not None
+            entry.entry_id in hass.data[DOMAIN][API_CLIENT]
+            and hass.data[DOMAIN][API_CLIENT][entry.entry_id] is not None
         ):
             try:
-                await hass.data[DOMAIN][WEBSOCKET_CLIENT][entry.entry_id].close()
+                await hass.data[DOMAIN][API_CLIENT][entry.entry_id].close()
             except:
                 pass
-        hass.data[DOMAIN][WEBSOCKET_CLIENT].pop(entry.entry_id)
+        hass.data[DOMAIN][API_CLIENT].pop(entry.entry_id)
 
     device_id = entry.data.get(CONF_DID)
     dispatchers = hass.data[DOMAIN][DISPATCHERS].pop(device_id)

custom_components/vaillant_plus/binary_sensor.py

@@ -15,8 +15,8 @@
 from homeassistant.helpers.entity import EntityCategory
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .client import VaillantDeviceApiClient
-from .const import CONF_DID, DISPATCHERS, DOMAIN, EVT_DEVICE_CONNECTED, WEBSOCKET_CLIENT
+from .client import VaillantClient
+from .const import CONF_DID, DISPATCHERS, DOMAIN, EVT_DEVICE_CONNECTED, API_CLIENT
 from .entity import VaillantEntity
 
 _LOGGER = logging.getLogger(__name__)
@@ -42,56 +42,56 @@ class VaillantBinarySensorDescription(
         name="Circulation",
         device_class=BinarySensorDeviceClass.RUNNING,
         entity_category=EntityCategory.DIAGNOSTIC,
-        on_state=True,
+        on_state=1,
     ),
     VaillantBinarySensorDescription(
         key="Heating_Enable",
         name="Heating",
         device_class=BinarySensorDeviceClass.RUNNING,
         entity_category=EntityCategory.DIAGNOSTIC,
-        on_state=True,
+        on_state=1,
     ),
     VaillantBinarySensorDescription(
         key="WarmStar_Tank_Loading_Enable",
         name="WarmStar tank loading",
         device_class=BinarySensorDeviceClass.RUNNING,
         entity_category=EntityCategory.DIAGNOSTIC,
-        on_state=True,
+        on_state=1,
     ),
     VaillantBinarySensorDescription(
         key="Enabled_Heating",
         name="Heating boiler",
         device_class=BinarySensorDeviceClass.RUNNING,
         entity_category=EntityCategory.DIAGNOSTIC,
-        on_state=True,
+        on_state=1,
     ),
     VaillantBinarySensorDescription(
         key="Enabled_DHW",
         name="Domestic hot water",
         device_class=BinarySensorDeviceClass.RUNNING,
         entity_category=EntityCategory.DIAGNOSTIC,
-        on_state=True,
+        on_state=1,
     ),
     VaillantBinarySensorDescription(
         key="BMU_Platform",
         name="BMU platform",
         # device_class=BinarySensorDeviceClass.RUNNING,
         entity_category=EntityCategory.DIAGNOSTIC,
-        on_state=True,
+        on_state=1,
     ),
     VaillantBinarySensorDescription(
         key="Weather_compensation",
         name="Weather compensation",
         device_class=BinarySensorDeviceClass.RUNNING,
         entity_category=EntityCategory.DIAGNOSTIC,
-        on_state=True,
+        on_state=1,
     ),
     VaillantBinarySensorDescription(
         key="RF_Status",
         name="EBus status",
         device_class=BinarySensorDeviceClass.CONNECTIVITY,
         entity_category=EntityCategory.DIAGNOSTIC,
-        on_state=True,
+        on_state=3,
     ),
     VaillantBinarySensorDescription(
         key="Boiler_info3_bit0",
@@ -115,7 +115,7 @@ async def async_setup_entry(
 ) -> bool:
     """Set up Vaillant binary sensors."""
     device_id = entry.data.get(CONF_DID)
-    client: VaillantDeviceApiClient = hass.data[DOMAIN][WEBSOCKET_CLIENT][
+    client: VaillantClient = hass.data[DOMAIN][API_CLIENT][
         entry.entry_id
     ]
 
@@ -154,7 +154,7 @@ class VaillantBinarySensorEntity(VaillantEntity, BinarySensorEntity):
 
     def __init__(
         self,
-        client: VaillantDeviceApiClient,
+        client: VaillantClient,
         description: VaillantBinarySensorDescription,
     ):
         super().__init__(client)
@@ -174,8 +174,10 @@ def update_from_latest_data(self, data: dict[str, Any]) -> None:
         if self.entity_description.key == "RF_Status":
             self._attr_is_on = value == 3
         elif self.entity_description.key == "Boiler_info3_bit0":
-            self._attr_is_on = value[0] == 1
+            self._attr_is_on = value.startswith("1")
         elif self.entity_description.key == "Boiler_info5_bit4":
-            self._attr_is_on = value[0] == 1
+            self._attr_is_on = value.startswith("1")
+        elif self.entity_description.on_state is not None:
+            self._attr_is_on = value == self.entity_description.on_state
         else:
             self._attr_is_on = value is True