diff --git a/pyhilo/api.py b/pyhilo/api.py
index 7683b44..267f8a6 100755
--- a/pyhilo/api.py
+++ b/pyhilo/api.py
@@ -363,8 +363,6 @@ 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)
 
         # Initialize WebsocketManager ic-dev21
         self.websocket_manager = WebsocketManager(
@@ -373,6 +371,8 @@ async def _async_post_init(self) -> None:
         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
+        # instanciate differently
         self.websocket = WebsocketClient(self.websocket_manager.devicehub)
         self.websocket2 = WebsocketClient(self.websocket_manager.challengehub)
 
@@ -381,6 +381,8 @@ async def refresh_ws_token(self) -> None:
         await self.websocket_manager.refresh_token(self.websocket_manager.devicehub)
         await self.websocket_manager.refresh_token(self.websocket_manager.challengehub)
 
+
+    # ic-dev21 not sure this is still needed? See websocket.py _async_negotiate
     async def post_devicehub_negociate(self) -> tuple[str, str]:
         LOG.debug("Getting websocket url")
         url = f"{AUTOMATION_DEVICEHUB_ENDPOINT}/negotiate"
diff --git a/pyhilo/websocket.py b/pyhilo/websocket.py
index 418ad59..b7b7180 100755
--- a/pyhilo/websocket.py
+++ b/pyhilo/websocket.py
@@ -417,9 +417,9 @@ def __init__(
         self.async_request = async_request
         self._state_yaml = state_yaml
         self._set_state = set_state_callback
-        self._shared_token = None  # ic-dev21 need to share the token
+        self._shared_token = None
 
-        # Initialize websocket configurations
+        # Initialize websocket configurations, more can be added here
         self.devicehub = WebsocketConfig(
             endpoint=AUTOMATION_DEVICEHUB_ENDPOINT, session=session
         )
@@ -431,7 +431,7 @@ async def initialize_websockets(self) -> None:
         """Initialize both websocket connections"""
         # ic-dev21 get token from device hub
         await self.refresh_token(self.devicehub, get_new_token=True)
-        # ic-dev21 reuse it for challenge hub
+        # ic-dev21 get token from challenge hub
         await self.refresh_token(self.challengehub, get_new_token=True)
 
     async def refresh_token(
@@ -446,7 +446,6 @@ async def refresh_token(
             config.url, self._shared_token = await self._negotiate(config)
             config.token = self._shared_token
         else:
-            # ic-dev21 reuse existing token but get new URL
             config.url, _ = await self._negotiate(config)
             config.token = self._shared_token
 
@@ -472,7 +471,7 @@ async def _negotiate(self, config: WebsocketConfig) -> Tuple[str, str]:
         # Save state
         state_key = (
             "websocket"
-            if config.endpoint == "AUTOMATION_DEVICEHUB_ENDPOINT"
+            if config.endpoint == AUTOMATION_DEVICEHUB_ENDPOINT
             else "websocket2"
         )
         await self._set_state(
@@ -521,7 +520,7 @@ async def _get_websocket_params(self, config: WebsocketConfig) -> None:
         # Save state
         state_key = (
             "websocket"
-            if config.endpoint == "AUTOMATION_DEVICEHUB_ENDPOINT"
+            if config.endpoint == AUTOMATION_DEVICEHUB_ENDPOINT
             else "websocket2"
         )
         LOG.debug(f"Calling set_state {state_key}_params")