Add support for setting HeaterMode via the Heater class

pyomnilogic_local/heater.py

@@ -40,7 +40,7 @@ class Heater(OmniEquipment[MSPVirtualHeater, TelemetryVirtualHeater]):
         min_temp: Minimum settable temperature (Fahrenheit)
 
     Properties (Telemetry):
-        mode: Current heater mode (OFF, HEAT, AUTO, etc.)
+        mode: Current heater mode (HEATING, COOLING, AUTO, or OFF.)
         current_set_point: Current target temperature (Fahrenheit)
         solar_set_point: Solar heater target temperature (Fahrenheit)
         enabled: Whether heater is enabled
@@ -243,6 +243,16 @@ async def set_temperature(self, temperature: int) -> None:
         # Always use Fahrenheit as that's what the OmniLogic system uses internally
         await self._api.async_set_heater(self.bow_id, self.system_id, temperature)
 
+    @control_method
+    async def set_mode(self, mode: HeaterMode) -> None:
+        """Set the heater operating mode: HEATING, COOLING, AUTO, or OFF."""
+
+        if self.bow_id is None or self.system_id is None:
+            msg = "Cannot set heater mode: bow_id or system_id is None"
+            raise OmniEquipmentNotInitializedError(msg)
+
+        await self._api.async_set_heater_mode(self.bow_id, self.system_id, mode)
+
     @control_method
     async def set_solar_temperature(self, temperature: int) -> None:
         """Set the solar heater set point.
@@ -272,3 +282,20 @@ async def set_solar_temperature(self, temperature: int) -> None:
 
         # Always use Fahrenheit as that's what the OmniLogic system uses internally
         await self._api.async_set_solar_heater(self.bow_id, self.system_id, temperature)
+
+    @property
+    def supports_cooling(self) -> bool:
+        """Return True if any physical heater equipment supports cooling."""
+        return any(
+            heater_equip.supports_cooling is True
+            for _, _, heater_equip in self.heater_equipment.items()
+        )
+
+    @property
+    def cooling_heater_equipment(self) -> tuple[HeaterEquipment, ...]:
+        """Return physical heater equipment that supports cooling."""
+        return tuple(
+            heater_equip
+            for _, _, heater_equip in self.heater_equipment.items()
+            if heater_equip.supports_cooling is True
+        )

tests/test_fixtures.py

@@ -13,13 +13,17 @@
 
 import json
 import pathlib
+from types import SimpleNamespace
 from typing import TYPE_CHECKING, Any
+from unittest.mock import AsyncMock, Mock
 
 import pytest
 
-from pyomnilogic_local.models.mspconfig import MSPConfig
-from pyomnilogic_local.models.telemetry import Telemetry
-from pyomnilogic_local.omnitypes import OmniType
+from pyomnilogic_local.collections import EquipmentDict
+from pyomnilogic_local.heater import Heater
+from pyomnilogic_local.models.mspconfig import MSPConfig, MSPVirtualHeater
+from pyomnilogic_local.models.telemetry import Telemetry, TelemetryHeater, TelemetryVirtualHeater
+from pyomnilogic_local.omnitypes import BackyardState, HeaterMode, HeaterState, HeaterWhyOn, OmniType
 
 if TYPE_CHECKING:
     from pyomnilogic_local._base import OmniEquipment
@@ -74,6 +78,112 @@ def get_equipment_by_type(msp: MSPConfig, omni_type: OmniType) -> list[Any]:
     return equipment
 
 
+def make_heater_equipment_data(system_id: int, name: str, *, supports_cooling: bool) -> dict[str, Any]:
+    """Create MSP config data for physical heater equipment."""
+    return {
+        "System-Id": system_id,
+        "Name": name,
+        "Type": "PET_HEATER",
+        "Heater-Type": "HTR_HEAT_PUMP",
+        "Enabled": True,
+        "Min-Speed-For-Operation": 45,
+        "Sensor-System-Id": 99,
+        "SupportsCooling": supports_cooling,
+    }
+
+
+def make_heater() -> Heater:
+    """Create a virtual heater with one cooling-capable and one heat-only unit."""
+    heater_config = MSPVirtualHeater.model_validate(
+        {
+            "System-Id": 20,
+            "Name": "Virtual Heater",
+            "Enabled": True,
+            "Current-Set-Point": 80,
+            "Max-Settable-Water-Temp": 104,
+            "Min-Settable-Water-Temp": 65,
+            "Operation": [
+                {OmniType.HEATER_EQUIP: make_heater_equipment_data(21, "Heat Pump", supports_cooling=True)},
+                {OmniType.HEATER_EQUIP: make_heater_equipment_data(22, "Gas Heater", supports_cooling=False)},
+            ],
+        },
+    )
+    heater_config.propagate_bow_id(7)
+
+    virtual_telemetry = TelemetryVirtualHeater.model_validate(
+        {
+            "@systemId": 20,
+            "@Current-Set-Point": 80,
+            "@enable": True,
+            "@SolarSetPoint": 0,
+            "@Mode": HeaterMode.HEATING,
+            "@SilentMode": 0,
+            "@whyHeaterIsOn": HeaterWhyOn.STOP_HEATER,
+        },
+    )
+    physical_telemetry = {
+        system_id: TelemetryHeater.model_validate(
+            {
+                "@systemId": system_id,
+                "@heaterState": HeaterState.OFF,
+                "@temp": 78,
+                "@enable": True,
+                "@priority": priority,
+                "@maintainFor": 24,
+            },
+        )
+        for priority, system_id in enumerate((21, 22), start=1)
+    }
+
+    def get_telem(system_id: int) -> TelemetryVirtualHeater | TelemetryHeater | None:
+        if system_id == 20:
+            return virtual_telemetry
+        return physical_telemetry.get(system_id)
+
+    telemetry = Mock(spec=Telemetry)
+    telemetry.get_telem_by_systemid = Mock(side_effect=get_telem)
+
+    omni = Mock()
+    omni._api = Mock()
+    omni._make_equipment_dict = Mock(side_effect=lambda items=None: EquipmentDict(items))
+    omni._telemetry_dirty = False
+    omni.backyard = SimpleNamespace(telemetry=SimpleNamespace(state=BackyardState.ON))
+
+    return Heater(omni, heater_config, telemetry)
+
+
+class TestHeaterControls:
+    """Tests for Heater behavior using fixture-style equipment data."""
+
+    @pytest.mark.asyncio
+    async def test_heater_set_mode(self) -> None:
+        """Test set_mode calls the heater mode API with the BoW and virtual heater IDs."""
+        heater = make_heater()
+        heater._api.async_set_heater_mode = AsyncMock()  # type: ignore[method-assign,union-attr]
+
+        await heater.set_mode(HeaterMode.COOLING)
+
+        heater._api.async_set_heater_mode.assert_called_once_with(7, 20, HeaterMode.COOLING)
+
+    def test_heater_supports_cooling_when_any_heater_equipment_supports_cooling(self) -> None:
+        """Test a virtual heater supports cooling when any physical heater equipment supports it."""
+        heater = make_heater()
+
+        assert len(heater.heater_equipment) == 2
+        assert heater.heater_equipment["Heat Pump"].supports_cooling is True
+        assert heater.heater_equipment["Gas Heater"].supports_cooling is False
+        assert heater.supports_cooling is True
+
+    def test_heater_cooling_equipment(self) -> None:
+        """Test cooling_equipment returns only physical heater equipment that supports cooling."""
+        heater = make_heater()
+
+        cooling_equipment = heater.cooling_heater_equipment
+
+        assert len(cooling_equipment) == 1
+        assert cooling_equipment[0].name == "Heat Pump"
+
+
 class TestIssue144:
     """Tests for issue-144.json fixture.