Merge pull request #236 from fersingb/Websocket

Added backward compatibility in the websocket branch

Author: ic-dev21

pyhilo/api.py

@@ -27,7 +27,7 @@
     API_NOTIFICATIONS_ENDPOINT,
     API_REGISTRATION_ENDPOINT,
     API_REGISTRATION_HEADERS,
-    AUTOMATION_DEVICEHUB_ENDPOINT,
+    AUTOMATION_CHALLENGE_ENDPOINT,
     DEFAULT_STATE_FILE,
     DEFAULT_USER_AGENT,
     FB_APP_ID,
@@ -51,7 +51,7 @@
     get_state,
     set_state,
 )
-from pyhilo.websocket import WebsocketClient
+from pyhilo.websocket import WebsocketClient, WebsocketManager
 
 
 class API:
@@ -81,9 +81,17 @@ def __init__(
         self.device_attributes = get_device_attributes()
         self.session: ClientSession = session
         self._oauth_session = oauth_session
+        self.websocket_devices: WebsocketClient
+        # Backward compatibility during transition to websocket for challenges. Currently the HA Hilo integration
+        # uses the .websocket attribute. Re-added this attribute and point to the same object as websocket_devices.
+        # Should be removed once the transition to the challenge websocket is completed everywhere.
         self.websocket: WebsocketClient
+        self.websocket_challenges: WebsocketClient
         self.log_traces = log_traces
         self._get_device_callbacks: list[Callable[..., Any]] = []
+        self.ws_url: str = ""
+        self.ws_token: str = ""
+        self.endpoint: str = ""
 
     @classmethod
     async def async_create(
@@ -132,6 +140,9 @@ async def async_get_access_token(self) -> str:
         if not self._oauth_session.valid_token:
             await self._oauth_session.async_ensure_token_valid()
 
+        access_token = str(self._oauth_session.token["access_token"])
+        LOG.debug(f"ic-dev21 access token is {access_token}")
+
         return str(self._oauth_session.token["access_token"])
 
     def dev_atts(
@@ -216,17 +227,24 @@ async def _async_request(
         :rtype: dict[str, Any]
         """
         kwargs.setdefault("headers", self.headers)
+        access_token = await self.async_get_access_token()
+
         if endpoint.startswith(API_REGISTRATION_ENDPOINT):
             kwargs["headers"] = {**kwargs["headers"], **API_REGISTRATION_HEADERS}
         if endpoint.startswith(FB_INSTALL_ENDPOINT):
             kwargs["headers"] = {**kwargs["headers"], **FB_INSTALL_HEADERS}
         if endpoint.startswith(ANDROID_CLIENT_ENDPOINT):
             kwargs["headers"] = {**kwargs["headers"], **ANDROID_CLIENT_HEADERS}
         if host == API_HOSTNAME:
-            access_token = await self.async_get_access_token()
             kwargs["headers"]["authorization"] = f"Bearer {access_token}"
         kwargs["headers"]["Host"] = host
 
+        # ic-dev21 trying Leicas suggestion
+        if endpoint.startswith(AUTOMATION_CHALLENGE_ENDPOINT):
+            # remove Ocp-Apim-Subscription-Key header to avoid 401 error
+            kwargs["headers"].pop("Ocp-Apim-Subscription-Key", None)
+            kwargs["headers"]["authorization"] = f"Bearer {access_token}"
+
         data: dict[str, Any] = {}
         url = parse.urljoin(f"https://{host}", endpoint)
         if self.log_traces:
@@ -303,8 +321,9 @@ async def _async_handle_on_backoff(self, _: dict[str, Any]) -> None:
                 LOG.info(
                     "401 detected on websocket, refreshing websocket token. Old url: {self.ws_url} Old Token: {self.ws_token}"
                 )
+                LOG.info(f"401 detected on {err.request_info.url}")
                 async with self._backoff_refresh_lock_ws:
-                    (self.ws_url, self.ws_token) = await self.post_devicehub_negociate()
+                    await self.refresh_ws_token()
                     await self.get_websocket_params()
                 return
 
@@ -354,30 +373,26 @@ async def _async_post_init(self) -> None:
         LOG.debug("Websocket postinit")
         await self._get_fid()
         await self._get_device_token()
-        await self.refresh_ws_token()
-        self.websocket = WebsocketClient(self)
 
-    async def refresh_ws_token(self) -> None:
-        (self.ws_url, self.ws_token) = await self.post_devicehub_negociate()
-        await self.get_websocket_params()
-
-    async def post_devicehub_negociate(self) -> tuple[str, str]:
-        LOG.debug("Getting websocket url")
-        url = f"{AUTOMATION_DEVICEHUB_ENDPOINT}/negotiate"
-        LOG.debug(f"devicehub URL is {url}")
-        resp = await self.async_request("post", url)
-        ws_url = resp.get("url")
-        ws_token = resp.get("accessToken")
-        LOG.debug("Calling set_state devicehub_negotiate")
-        await set_state(
-            self._state_yaml,
-            "websocket",
-            {
-                "url": ws_url,
-                "token": ws_token,
-            },
+        # Initialize WebsocketManager ic-dev21
+        self.websocket_manager = WebsocketManager(
+            self.session, self.async_request, self._state_yaml, set_state
         )
-        return (ws_url, ws_token)
+        await self.websocket_manager.initialize_websockets()
+
+        # Create both websocket clients
+        # ic-dev21 need to work on this as it can't lint as is, may need to
+        # instantiate differently
+        self.websocket_devices = WebsocketClient(self.websocket_manager.devicehub)
+
+        # For backward compatibility during the transition to challengehub websocket 
+        self.websocket = self.websocket_devices
+        self.websocket_challenges = WebsocketClient(self.websocket_manager.challengehub)
+
+    async def refresh_ws_token(self) -> None:
+        """Refresh the websocket token."""
+        await self.websocket_manager.refresh_token(self.websocket_manager.devicehub)
+        await self.websocket_manager.refresh_token(self.websocket_manager.challengehub)
 
     async def get_websocket_params(self) -> None:
         uri = parse.urlparse(self.ws_url)

pyhilo/const.py

@@ -42,6 +42,8 @@
 
 # Automation server constant
 AUTOMATION_DEVICEHUB_ENDPOINT: Final = "/DeviceHub"
+AUTOMATION_CHALLENGE_ENDPOINT: Final = "/ChallengeHub"
+
 
 # Request constants
 DEFAULT_USER_AGENT: Final = f"PyHilo/{PYHILO_VERSION} HomeAssistant/{homeassistant.core.__version__} aiohttp/{aiohttp.__version__} Python/{platform.python_version()}"

pyhilo/event.py

@@ -1,10 +1,13 @@
 """Event object """
 from datetime import datetime, timedelta, timezone
+import logging
 import re
 from typing import Any, cast
 
 from pyhilo.util import camel_to_snake, from_utc_timestamp
 
+LOG = logging.getLogger(__package__)
+
 
 class Event:
     setting_deadline: datetime
@@ -126,9 +129,12 @@ def current_phase_times(self) -> dict[str, datetime]:
     @property
     def state(self) -> str:
         now = datetime.now(self.preheat_start.tzinfo)
-        if self.pre_cold_start <= now < self.pre_cold_end:
+        if self.pre_cold_start and self.pre_cold_start <= now < self.pre_cold_end:
             return "pre_cold"
-        elif self.appreciation_start <= now < self.appreciation_end:
+        elif (
+            self.appreciation_start
+            and self.appreciation_start <= now < self.appreciation_end
+        ):
             return "appreciation"
         elif self.preheat_start > now:
             return "scheduled"
@@ -138,9 +144,12 @@ def state(self) -> str:
             return "reduction"
         elif self.recovery_start <= now < self.recovery_end:
             return "recovery"
+        elif now >= self.recovery_end + timedelta(minutes=5):
+            return "off"
         elif now >= self.recovery_end:
             return "completed"
         elif self.progress:
             return self.progress
+
         else:
             return "unknown"

pyhilo/util/__init__.py

@@ -1,6 +1,7 @@
 """Define utility modules."""
 import asyncio
 from datetime import datetime, timedelta
+import logging
 import re
 from typing import Any, Callable
 
@@ -9,6 +10,9 @@
 
 from pyhilo.const import LOG  # noqa: F401
 
+LOG = logging.getLogger(__package__)
+
+
 CAMEL_REX_1 = re.compile("(.)([A-Z][a-z]+)")
 CAMEL_REX_2 = re.compile("([a-z0-9])([A-Z])")
 
@@ -35,7 +39,11 @@ def snake_to_camel(string: str) -> str:
 def from_utc_timestamp(date_string: str) -> datetime:
     from_zone = tz.tzutc()
     to_zone = tz.tzlocal()
-    return parse(date_string).replace(tzinfo=from_zone).astimezone(to_zone)
+    dt = parse(date_string)
+    if dt.tzinfo is None:  # Only replace tzinfo if not already set
+        dt = dt.replace(tzinfo=from_zone)
+    output = dt.astimezone(to_zone)
+    return output
 
 
 def time_diff(ts1: datetime, ts2: datetime) -> timedelta: