diff --git a/pyhilo/api.py b/pyhilo/api.py
index 7683b44..9191923 100755
--- a/pyhilo/api.py
+++ b/pyhilo/api.py
@@ -363,7 +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
diff --git a/pyhilo/websocket.py b/pyhilo/websocket.py
index 418ad59..96b8aa1 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
 
@@ -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")