Keeping the right tokens for every websocket connection

Author: Leicas

pyhilo/api.py

@@ -311,8 +311,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
 
@@ -362,7 +363,7 @@ 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()
+        # await self.refresh_ws_token()
         # self.websocket = WebsocketClient(self)
 
         # Initialize WebsocketManager ic-dev21
@@ -376,8 +377,9 @@ async def _async_post_init(self) -> None:
         self.websocket2 = WebsocketClient(self, self.websocket_manager.challengehub)
 
     async def refresh_ws_token(self) -> None:
-        (self.ws_url, self.ws_token) = await self.post_devicehub_negociate()
-        await self.get_websocket_params()
+        """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 post_devicehub_negociate(self) -> tuple[str, str]:
         LOG.debug("Getting websocket url")

pyhilo/websocket.py

@@ -424,7 +424,7 @@ async def initialize_websockets(self) -> None:
         # ic-dev21 get token from device hub
         await self.refresh_token(self.devicehub, get_new_token=True)
         # ic-dev21 reuse it for challenge hub
-        await self.refresh_token(self.challengehub, get_new_token=False)
+        await self.refresh_token(self.challengehub, get_new_token=True)
 
     async def refresh_token(
         self, config: WebsocketConfig, get_new_token: bool = True
@@ -459,11 +459,8 @@ async def _negotiate(self, config: WebsocketConfig) -> Tuple[str, str]:
 
         resp = await self.async_request("post", url)
         ws_url = resp.get("url")
-        ws_token = (
-            resp.get("accessToken")
-            if self._shared_token is None
-            else self._shared_token
-        )
+        ws_token =   resp.get("accessToken")
+          
 
         # Save state
         state_key = (