add Heater metrics if Heater is enabled

custom_components/astralpool_halo_chlorinator/binary_sensor.py

@@ -56,6 +56,15 @@
     # ),
 }
 
+HEATER_BINARY_SENSOR_TYPES: dict[str, BinarySensorEntityDescription] = {
+    "HeaterOn": BinarySensorEntityDescription(
+        key="HeaterOn",
+        device_class=BinarySensorDeviceClass.RUNNING,
+        icon="mdi:fire",
+        name="Heater On",
+    )
+}
+
 
 async def async_setup_entry(
     hass: HomeAssistant,
@@ -64,6 +73,15 @@ async def async_setup_entry(
 ) -> None:
     """Set up Chlorinator binary sensors from a config entry."""
     data: ChlorinatorData = hass.data[DOMAIN][entry.entry_id]
+    coordinator = data.coordinator
+
+    async def add_heater_binary_sensor_callback():
+        await add_heater_binary_sensors(
+            hass, coordinator.added_entities, async_add_entities, coordinator
+        )
+
+    coordinator.add_heater_binary_sensor_callback = add_heater_binary_sensor_callback
+    await coordinator.async_config_entry_first_refresh()
 
     entities = [
         ChlorinatorBinarySensor(data.coordinator, sensor_desc)
@@ -72,6 +90,22 @@ async def async_setup_entry(
     async_add_entities(entities)
 
 
+async def add_heater_binary_sensors(
+    hass, added_entities, async_add_entities, coordinator
+):
+    """Setup Heater sensors if enabled"""
+
+    new_entities = []
+    for sensor_type, sensor_desc in HEATER_BINARY_SENSOR_TYPES.items():
+        unique_id = f"hchlor_{sensor_type}".lower()
+        if unique_id not in added_entities:
+            new_entities.append(HeaterBinarySensor(coordinator, sensor_desc))
+            added_entities.add(unique_id)
+
+    if new_entities:
+        async_add_entities(new_entities)
+
+
 class ChlorinatorBinarySensor(CoordinatorEntity, BinarySensorEntity):
     """Representation of a Clorinator binary sensor."""
 
@@ -103,3 +137,32 @@ def device_info(self) -> DeviceInfo | None:
     def is_on(self) -> bool:
         """Return the state of the sensor."""
         return self.coordinator.data.get(self._sensor)
+
+
+class HeaterBinarySensor(CoordinatorEntity, BinarySensorEntity):
+    """Representation of a Clorinator binary sensor."""
+
+    _attr_name = "Pump is operating"
+
+    def __init__(self, coordinator, sensor_desc: BinarySensorEntityDescription):
+        """Initialize the sensor."""
+        super().__init__(coordinator)
+        self._sensor = sensor_desc.key
+        self._attr_unique_id = f"hchlor_{self._sensor}".lower()
+        self.entity_description = sensor_desc
+        self._attr_name = sensor_desc.name
+        self._attr_device_class = sensor_desc.device_class
+
+    @property
+    def device_info(self) -> DeviceInfo | None:
+        return {
+            "identifiers": {(DOMAIN, "HCHLOR")},
+            "name": "HCHLOR",
+            "model": "Halo Chlor",
+            "manufacturer": "Astral Pool",
+        }
+
+    @property
+    def is_on(self) -> bool:
+        """Return the state of the sensor."""
+        return self.coordinator.data.get(self._sensor)

custom_components/astralpool_halo_chlorinator/coordinator.py

@@ -15,7 +15,7 @@
 _LOGGER = logging.getLogger(__name__)
 
 
-class ChlorinatorDataUpdateCoordinator(DataUpdateCoordinator[dict[str, Any]]):
+class ChlorinatorDataUpdateCoordinator(DataUpdateCoordinator):
     """Data coordinator for getting Chlorinator updates."""
 
     def __init__(self, hass: HomeAssistant, chlorinator: HaloChlorinatorAPI) -> None:
@@ -34,6 +34,9 @@ def __init__(self, hass: HomeAssistant, chlorinator: HaloChlorinatorAPI) -> None
             manufacturer="Astral Pool",
             name="HCHLOR",
         )
+        self.added_entities = set()
+        self.add_heater_sensor_callback = None
+        self.add_heater_binary_sensor_callback = None
 
     async def _async_update_data(self):
         """Fetch data from API endpoint."""
@@ -42,15 +45,28 @@ async def _async_update_data(self):
         if self._data_age >= 3:  # 3 polling events = 60 seconds
             try:
                 data = await self.chlorinator.async_gatherdata()
-                _LOGGER.debug("halo_ble_client finish: %s", data)
+                _LOGGER.debug("halo_ble_client finish: %s", dict(sorted(data.items())))
             except Exception as e:
                 _LOGGER.warning("Failed _gatherdata: %s %s", self._data_age, e)
                 data = {}
             if data != {}:
                 self.data = data
                 self._data_age = 0
+
+                if "HeaterEnabled" in data and data["HeaterEnabled"] == 1:
+                    _LOGGER.debug("HeaterEnabled : %s", data["HeaterEnabled"])
+                    if (
+                        self.add_heater_sensor_callback
+                        and self.add_heater_binary_sensor_callback is not None
+                    ):
+                        await self.add_heater_sensor_callback()
+                        await self.add_heater_binary_sensor_callback()
+                    else:
+                        _LOGGER.warning("add_heater_callback(s) not set")
+
             elif self._data_age >= 15:  # 15 polling events  = 5 minutes
                 self.data = {}
                 _LOGGER.error("Failed _gatherdata, giving up: %s", self._data_age)
                 raise UpdateFailed("Error communicating with API")
+
         return self.data

custom_components/astralpool_halo_chlorinator/sensor.py

@@ -152,6 +152,16 @@
     ),
 }
 
+HEATER_SENSOR_TYPES: dict[str, SensorEntityDescription] = {
+    "HeaterMode": SensorEntityDescription(
+        key="HeaterMode",
+        icon="mdi:heat-pump",
+        name="Heater Mode",
+        native_unit_of_measurement=None,
+        device_class=SensorDeviceClass.ENUM,
+    )
+}
+
 
 async def async_setup_entry(
     hass: HomeAssistant,
@@ -160,6 +170,16 @@ async def async_setup_entry(
 ) -> None:
     """Set up Chlorinator from a config entry."""
     data: ChlorinatorData = hass.data[DOMAIN][entry.entry_id]
+    coordinator = data.coordinator
+
+    async def add_heater_sensor_callback():
+        await add_heater_sensors(
+            hass, coordinator.added_entities, async_add_entities, coordinator
+        )
+
+    coordinator.add_heater_sensor_callback = add_heater_sensor_callback
+    await coordinator.async_config_entry_first_refresh()
+
     entities = [
         ChlorinatorSensor(data.coordinator, sensor_desc)
         for sensor_desc in CHLORINATOR_SENSOR_TYPES
@@ -168,6 +188,20 @@ async def async_setup_entry(
     async_add_entities(entities)
 
 
+async def add_heater_sensors(hass, added_entities, async_add_entities, coordinator):
+    """Setup Heater sensors if enabled"""
+
+    new_entities = []
+    for sensor_type, sensor_desc in HEATER_SENSOR_TYPES.items():
+        unique_id = f"hchlor_{sensor_type}".lower()
+        if unique_id not in added_entities:
+            new_entities.append(HeaterSensor(coordinator, sensor_desc))
+            added_entities.add(unique_id)
+
+    if new_entities:
+        async_add_entities(new_entities)
+
+
 class ChlorinatorSensor(
     CoordinatorEntity[ChlorinatorDataUpdateCoordinator], SensorEntity
 ):
@@ -203,3 +237,31 @@ def device_info(self) -> DeviceInfo | None:
     @property
     def native_value(self):
         return self.coordinator.data.get(self._sensor)
+
+
+class HeaterSensor(CoordinatorEntity[ChlorinatorDataUpdateCoordinator], SensorEntity):
+    """Representation of a Heater Sensor."""
+
+    def __init__(self, coordinator, sensor_desc: SensorEntityDescription) -> None:
+        """Initialize the sensor."""
+        super().__init__(coordinator)
+        self._sensor = sensor_desc.key
+        self._attr_unique_id = f"hchlor_{self._sensor}".lower()
+        self.entity_description = sensor_desc
+        self._attr_name = sensor_desc.name
+        self._attr_native_unit_of_measurement = sensor_desc.native_unit_of_measurement
+
+    @property
+    def device_info(self) -> DeviceInfo | None:
+        # Device info remains the same
+        return {
+            "identifiers": {(DOMAIN, "HCHLOR")},
+            "name": "HCHLOR",
+            "model": "Halo Chlor",
+            "manufacturer": "Astral Pool",
+        }
+
+    @property
+    def native_value(self):
+        # Use self._sensor to fetch the relevant data from coordinator
+        return self.coordinator.data.get(self._sensor)