Reconnect immediately if connection fails on token expiry

Author: steventux

src/relay_listener.py

@@ -16,6 +16,7 @@
 
 from dotenv import load_dotenv
 from websockets.asyncio.client import connect
+from websockets.exceptions import ConnectionClosedError
 
 from services.mwl.create_worklist_item import CreateWorklistItem
 from services.storage import MWLStorage
@@ -54,7 +55,7 @@ async def listen(self):
 
         logger.info(f"Connecting to Azure Relay: {self.relay_uri.hybrid_connection_name}...")
 
-        async with connect(self.relay_uri.connection_url(), compression=None) as websocket:
+        async with self._connect() as websocket:
             logger.info("Connected - waiting for worklist actions...")
 
             async for message in websocket:
@@ -88,6 +89,10 @@ def process_action(self, payload: dict):
 
         return action_class(self.storage).call(payload)
 
+    def _connect(self):
+        """Connect to Azure Relay."""
+        return connect(self.relay_uri.connection_url(), compression=None)
+
 
 class RelayURI:
     def __init__(self):
@@ -133,6 +138,12 @@ async def main():
         except KeyboardInterrupt:
             logger.warning("\nShutting down...")
             break
+        except ConnectionClosedError as e:
+            if "ExpiredToken" in str(e) and e.code == 1011:
+                logger.info("SAS token expired, refreshing...")
+            else:
+                logger.warning(f"Connection closed with code {e.code}: {e.reason}")
+                raise e
         except Exception as e:
             logger.warning(f"Connection error: {e}")
             logger.warning("Retrying in 5 seconds...")

tests/test_relay_listener.py

@@ -1,9 +1,11 @@
 import json
-from unittest.mock import patch
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
+from websockets.exceptions import ConnectionClosedError
+from websockets.frames import Close, CloseCode
 
-from relay_listener import RelayListener, RelayURI
+from relay_listener import RelayListener, RelayURI, main
 from services.storage import WorklistItem
 
 
@@ -121,3 +123,25 @@ def test_relay_uri_connection_url(self):
             "&sb-hc-token=SharedAccessSignature+sr%3Dhttp%253A%252F%252Ftest-namespace"
             "%252Ftest-connection%26sig%3DPMcelSnwGlYX2xFo9Y2aGCg%252BvJ6LsHujiRrA1L6VnP0%253D%26se%3D1003600%26skn%3Dtest-key-name"
         )
+
+
+@patch("relay_listener.logger", new_callable=MagicMock)
+@patch("relay_listener.MWLStorage", new_callable=MagicMock)
+@patch("relay_listener.RelayListener")
+@patch("asyncio.sleep", new_callable=AsyncMock)
+@pytest.mark.asyncio
+async def test_main_handles_connection_closed_and_keyboard_interrupt(mock_sleep, mock_relay_listener, mock_mwl_storage, mock_logger):
+    relay_listener_instance = mock_relay_listener.return_value
+    relay_listener_instance.listen = AsyncMock()
+
+    relay_listener_instance.listen.side_effect = [
+        ConnectionClosedError(Close(CloseCode.INTERNAL_ERROR, "ExpiredToken"), None),
+        KeyboardInterrupt()
+    ]
+
+    await main()
+
+    assert relay_listener_instance.listen.call_count == 2
+    mock_logger.info.assert_any_call("Socket Listener Starting...")
+    mock_logger.info.assert_any_call("SAS token expired, refreshing...")
+    mock_logger.warning.assert_any_call("\nShutting down...")