implemnt that coct app usages nextstage instead of updating stage on …

…endpoint

custom_components/coct_loadshedding/coct_interface.py

@@ -1,4 +1,5 @@
 import ssl
+import datetime
 
 from aiohttp.client_exceptions import ClientConnectorError, ServerDisconnectedError
 from aiohttp_retry import RetryClient
@@ -11,14 +12,15 @@ class coct_interface:
     def __init__(self):
         """Initializes class parameters"""
 
-        self.base_url = "https://d42sspn7yra3u.cloudfront.net"
+        self.base_url_eskom = "https://loadshedding.eskom.co.za/LoadShedding"
+        self.base_url_ct = "https://d42sspn7yra3u.cloudfront.net"
         self.headers = {
             "user_agent": "Mozilla/5.0 (X11; Linux x86_64; rv:69.0) Gecko/20100101 Firefox/69.0"
         }
         self.ssl_context = ssl.create_default_context()
         self.ssl_context.set_ciphers("DEFAULT@SECLEVEL=1")
 
-    async def async_query_api(self, endpoint, payload=None):
+    async def async_query_api(self, base, endpoint, payload=None):
         """Queries a given endpoint on the CoCT loadshedding API with the specified payload
 
         Args:
@@ -31,21 +33,55 @@ async def async_query_api(self, endpoint, payload=None):
         async with RetryClient() as client:
             # The CoCT API occasionally drops incoming connections, implement reies
             async with client.get(
-                url=self.base_url + endpoint,
+                url=base + endpoint,
                 headers=self.headers,
                 params=payload,
                 ssl=self.ssl_context,
                 retry_attempts=50,
                 retry_exceptions={
-                    ClientConnectorError,
-                    ServerDisconnectedError,
-                    ConnectionError,
-                    OSError,
+                   ClientConnectorError,
+                   ServerDisconnectedError,
+                   ConnectionError,
+                   OSError,
                 },
             ) as res:
                 return await res.json()
 
-    async def async_get_stage(self, attempts=5):
+    async def async_get_stage_eskom(self, attempts=5):
+        """Fetches the current loadshedding stage from the Eskom API
+        Args:
+            attempts (int, optional): The number of attempts to query a sane value from the Eskom API. Defaults to 5.
+        Returns:
+            The loadshedding stage if the query succeeded, else `None`
+        """
+
+        # Placeholder for returned loadshedding stage
+        api_result = None
+
+        # Query the API until a sensible (> 0) value is received, or the number of attempts is exceeded
+        for attempt in range(attempts):
+            res = await self.async_query_api(self.base_url_eskom, "/GetStatus")
+
+            # Check if the API returned a valid response
+            if res:
+                # Store the response
+                api_result = res
+
+                # Only return the result if the API returned a non-negative stage, otherwise retry
+                if int(res) > 0:
+                    # Return the current loadshedding stage by subtracting 1 from the query result
+                    return int(res) - 1
+
+        if api_result:
+            # If the API is up but returning "invalid" stages (< 0), simply return 0
+            return 0
+        else:
+            # If the API the query did not succeed after the number of attempts has been exceeded, raise an exception
+            raise Exception(
+                f"Error, no response received from API after {attempts} attempts"
+            )
+
+    async def async_get_stage_coct(self, attempts=5):
         """Fetches the current loadshedding stage from the CoCT API
 
         Args:
@@ -60,7 +96,7 @@ async def async_get_stage(self, attempts=5):
 
         # Query the API until a sensible (> 0) value is received, or the number of attempts is exceeded
         for attempt in range(attempts):
-            res = await self.async_query_api("?")
+            res = await self.async_query_api(self.base_url_ct, "?")
 
             # Check if the API returned a valid response
             if res:
@@ -84,24 +120,44 @@ async def async_get_stage(self, attempts=5):
 
     async def async_get_data(self, coct_area):
         """Fetches data from the loadshedding API"""
-        json = await self.async_get_stage()
+        d = datetime.datetime.now()
+        #d = datetime.datetime.strptime("2022-07-10T16:20", '%Y-%m-%dT%H:%M')
+        load_shedding_active = False
+        next_load_shedding_slot = "N/A"
+
+        # grab json and stage
+        json = await self.async_get_stage_coct()
+        stage_eskom = await self.async_get_stage_eskom()
         stage = json[0]['currentStage']
-        if stage > 0:
-            load_shedding_active = isLoadSheddingNow(stage, coct_area)["status"]
-            next_load_shedding_slot = getNextTimeSlot(stage, coct_area)["date"]
-        else:
-            load_shedding_active = False
-            next_load_shedding_slot = "N/A"
+        next_stage = json[0]['nextStage']
+        next_stage_start_time = datetime.datetime.strptime(json[0]['nextStageStartTime'], '%Y-%m-%dT%H:%M')
+        last_updated = datetime.datetime.strptime(json[0]['lastUpdated'], '%Y-%m-%dT%H:%M:%S.000Z')
+
+        # if loadshedding active calculate slots for area if area set
+        if stage > 0 and coct_area > 0:
+            try:
+                next_load_shedding_slot = getNextTimeSlot(stage, coct_area)["date"]
+                load_shedding_active = isLoadSheddingNow(stage, coct_area)["status"]
+            except:
+                pass
+
+        # CoCT app works out different 'stage' if after 'next_stage_start_time'
+        if next_stage_start_time < d:
+            stage = next_stage
+        # Just in case, check eskom stage, if it is lower than stage use that
+        if stage_eskom < stage:
+            stage = stage_eskom
 
         data = {
             "data": {
                 "stage": stage,
+                "stage_eskom": stage_eskom,
                 "load_shedding_active": load_shedding_active,
                 "coct_area": coct_area,
                 "next_load_shedding_slot": next_load_shedding_slot,
-                "next_stage": json[0]['nextStage'],
-                "next_stage_start_time": json[0]['nextStageStartTime'],
-                "last_updated": json[0]['lastUpdated']
+                "next_stage": next_stage,
+                "next_stage_start_time": next_stage_start_time,
+                "last_updated": last_updated
             },
         }
         return data