From ff06d6110ca02a87a070f7bfca24c6322534899d Mon Sep 17 00:00:00 2001
From: vectorcatch <vectorcatch@users.noreply.github.com>
Date: Fri, 11 Mar 2022 05:16:30 -0500
Subject: [PATCH] Add set_limit and helpers for enabling and disabling (#16)

---
 pyjuicenet/__init__.py | 58 ++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 58 insertions(+)

diff --git a/pyjuicenet/__init__.py b/pyjuicenet/__init__.py
index 9f82f66..429b452 100644
--- a/pyjuicenet/__init__.py
+++ b/pyjuicenet/__init__.py
@@ -17,6 +17,7 @@ class Charger:
     def __init__(self, json_settings, api):
         """Create a Charger."""
         self.json_settings = json_settings
+        self.json_info = {}
         self.json_state = {}
         self.api = api
         self.last_updated_at = 0
@@ -45,6 +46,13 @@ async def update_state(self, force=False) -> bool:
         self.json_state = json_state
         return json_state["success"]
 
+    async def update_info(self) -> bool:
+        """Update device info with latest info from API."""
+        self.last_updated_at = time.time()
+        json_info = await self.api.get_info(self)
+        self.json_info = json_info
+        return json_info["success"]
+
     @property
     def voltage(self) -> int:
         """Get the voltage."""
@@ -85,6 +93,40 @@ def override_time(self) -> int:
         """Get the override time."""
         return self.json_state.get("override_time")
 
+    @property
+    def max_charging_amperage(self) -> int:
+        """Get the maximum charging limit time from the smaller of the wire rating and unit rating.
+        This can be used along with set limit to enable the charger at full capacity"""
+        return min(self.json_info.get("amps_wire_rating"), self.json_info.get("amps_unit_rating"))
+
+    @property
+    def current_charging_amperage_limit(self) -> int:
+        """Get the current amperage charging limit for the charger."""
+        return self.json_state.get("charging", {}).get("amps_limit")
+
+    @property
+    def charger_enabled(self) -> bool:
+        """Is the charger allowing amps to flow."""
+        return self.current_charging_amperage_limit > 0
+
+    async def set_charging_amperage_limit(self, amperage: int) -> bool:
+        """Set the amperage limit of the charger. 0 will disable the charger."""
+        response = await self.api.set_limit(self, amperage)
+
+        # Update state so that the amperage limit is correctly reflected
+        if response['success']:
+            await self.update_state(True)
+
+        return response['success']
+
+    async def enable_charger(self):
+        """Enable charger for max support amperage"""
+        return await self.set_charging_amperage_limit(self.max_charging_amperage)
+
+    async def disable_charger(self):
+        """Disable all charging activity"""
+        return await self.set_charging_amperage_limit(0)
+
     async def set_override(self, charge_now) -> bool:
         """Set to override schedule or not."""
         override_time = 0
@@ -164,6 +206,7 @@ async def get_devices(self):
         for unit in units_json:
             device = Charger(unit, self)
             await device.update_state()
+            await device.update_info()
             devices.append(device)
 
         return devices
@@ -196,6 +239,21 @@ async def get_info(self, charger: Charger):
         )
         return await response.json()
 
+    async def set_limit(self, charger: Charger, amperage_limit: int):
+        """Set the amperage limit of the charger. 0 will disable the charger."""
+        data = {
+            "cmd": "set_limit",
+            "device_id": self.uuid,
+            "token": charger.token,
+            "amperage": amperage_limit,
+            "account_token": self.api_token
+        }
+
+        response = await self.session.post(
+            f"{BASE_URL}/box_api_secure", json=data,
+        )
+        return await response.json()
+
     async def set_override(
         self,
         charger: Charger,