Add support for batteries API (#533)

* Add support for batteries API

* Update homewizard_energy/models.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Fix type discrepancy

* Fix duplicate snapshot

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: DCSBL

homewizard_energy/homewizard_energy.py

@@ -9,7 +9,7 @@
 
 from .const import LOGGER
 from .errors import UnsupportedError
-from .models import CombinedModels, Device, Measurement, State, System
+from .models import Batteries, CombinedModels, Device, Measurement, State, System
 
 
 class HomeWizardEnergy:
@@ -63,15 +63,20 @@ async def fetch_data(coroutine):
             except (UnsupportedError, NotImplementedError):
                 return None
 
-        device, measurement, system, state = await asyncio.gather(
+        device, measurement, system, state, batteries = await asyncio.gather(
             fetch_data(self.device()),
             fetch_data(self.measurement()),
             fetch_data(self.system()),
             fetch_data(self.state()),
+            fetch_data(self.batteries()),
         )
 
         return CombinedModels(
-            device=device, measurement=measurement, system=system, state=state
+            device=device,
+            measurement=measurement,
+            system=system,
+            state=state,
+            batteries=batteries,
         )
 
     async def device(self, reset_cache: bool = False) -> Device:
@@ -104,6 +109,10 @@ async def state(
         """Get/set the state."""
         raise UnsupportedError("State is not supported")
 
+    async def batteries(self, mode: Batteries.Mode | None = None) -> Batteries:
+        """Get the batteries."""
+        raise UnsupportedError("Batteries are not supported")
+
     async def identify(
         self,
     ) -> None:

homewizard_energy/models.py

@@ -22,7 +22,7 @@ class AwesomeVersionSerializationStrategy(SerializationStrategy, use_annotations
 
     def serialize(self, value: AwesomeVersion) -> str:
         """Serialize AwesomeVersion object to string."""
-        return str(value)
+        return str(value)  # pragma: no cover
 
     def deserialize(self, value: str) -> AwesomeVersion | None:
         """Deserialize string to AwesomeVersion object."""
@@ -65,18 +65,23 @@ class CombinedModels:
     measurement: Measurement
     state: State | None
     system: System | None
+    batteries: Batteries | None
 
+    # pylint: disable=too-many-arguments
+    # pylint: disable=too-many-positional-arguments
     def __init__(
         self,
         device: Device,
         measurement: Measurement,
         state: State | None,
         system: System | None,
+        batteries: Batteries | None = None,
     ):
         self.device = device
         self.measurement = measurement
         self.state = state
         self.system = system
+        self.batteries = batteries
 
         # Move things around for backwards compatibility
         ## measurement.wifi_ssid -> system.wifi_ssid
@@ -183,6 +188,10 @@ def supports_telegram(self) -> bool:
         """Return if the device supports telegram."""
         return self.product_type == Model.P1_METER
 
+    def supports_batteries(self) -> bool:
+        """Return if the device supports batteries."""
+        return self.product_type == Model.P1_METER
+
 
 @dataclass(kw_only=True)
 class Measurement(BaseModel):
@@ -601,6 +610,36 @@ def __post_deserialize__(cls, obj: System) -> System:
         return obj
 
 
+@dataclass
+class BatteriesUpdate(UpdateBaseModel):
+    """Represent Batteries update config."""
+
+    mode: Batteries.Mode | None = field(default=None)
+
+
+@dataclass(kw_only=True)
+class Batteries(BaseModel):
+    """Represent Batteries config."""
+
+    class Mode(StrEnum):
+        """Device type allocations."""
+
+        ZERO = "zero"
+        TO_FULL = "to_full"
+        STANDBY = "standby"
+
+    mode: Mode | None = field(
+        default=None,
+        metadata={
+            "deserialize": lambda x: Batteries.Mode.__members__.get(x.upper(), None)
+        },
+    )
+    power_w: float | None = field(default=None)
+    target_power_w: float | None = field(default=None)
+    max_consumption_w: float | None = field(default=None)
+    max_production_w: float | None = field(default=None)
+
+
 @dataclass(kw_only=True)
 class Token(BaseModel):
     """Represent Token."""

homewizard_energy/v2/__init__.py

@@ -19,12 +19,22 @@
 from ..errors import (
     DisabledError,
     InvalidUserNameError,
+    NotFoundError,
     RequestError,
     ResponseError,
     UnauthorizedError,
+    UnsupportedError,
 )
 from ..homewizard_energy import HomeWizardEnergy
-from ..models import Device, Measurement, System, SystemUpdate, Token
+from ..models import (
+    Batteries,
+    BatteriesUpdate,
+    Device,
+    Measurement,
+    System,
+    SystemUpdate,
+    Token,
+)
 from .cacert import CACERT
 
 T = TypeVar("T")
@@ -137,6 +147,36 @@ async def system(
         system = System.from_json(response)
         return system
 
+    @authorized_method
+    async def batteries(
+        self,
+        mode: Batteries.Mode | None = None,
+    ) -> Batteries:
+        """Return the batteries object."""
+
+        if self._device is not None and self._device.supports_batteries() is False:
+            raise UnsupportedError("Batteries is not supported")
+
+        try:
+            if mode is not None:
+                data = BatteriesUpdate(
+                    mode=mode,
+                ).to_dict()
+                status, response = await self._request(
+                    "/api/batteries", method=METH_PUT, data=data
+                )
+            else:
+                status, response = await self._request("/api/batteries")
+
+            if status != HTTPStatus.OK:
+                error = json.loads(response).get("error", response)
+                raise RequestError(f"Failed to get batteries: {error}")
+        except NotFoundError as exception:
+            # The batteries endpoint is not available on the device
+            raise UnsupportedError("Batteries is not supported") from exception
+
+        return Batteries.from_json(response)
+
     @authorized_method
     async def identify(
         self,
@@ -275,6 +315,8 @@ async def _request(
             case HTTPStatus.NO_CONTENT:
                 # No content, just return
                 return (HTTPStatus.NO_CONTENT, None)
+            case HTTPStatus.NOT_FOUND:
+                raise NotFoundError("Resource not found")
             case HTTPStatus.OK:
                 pass
 

tests/test_homewizard_energy.py

@@ -18,6 +18,7 @@
         ("state", UnsupportedError),
         ("identify", NotImplementedError),
         ("reboot", UnsupportedError),
+        ("batteries", UnsupportedError),
     ],
 )
 async def test_base_class_raises_notimplementederror(