Improve error handling for connection closed errors

If the connection closes for some other reason than token expiry, retry in 5 seconds.

Author: steventux

src/relay_listener.py

@@ -32,6 +32,7 @@
 ACTIONS = {
     "worklist.create_item": CreateWorklistItem,
 }
+EXPIRED_TOKEN = "ExpiredToken"
 
 
 class RelayListener:
@@ -140,11 +141,15 @@ async def main():
             logger.warning("\nShutting down...")
             break
         except ConnectionClosedError as e:
-            if "ExpiredToken" in str(e) and e.code == CloseCode.INTERNAL_ERROR.value:
+            code = e.rcvd.code if e.rcvd else "N/A"
+            reason = e.rcvd.reason if e.rcvd else "N/A"
+
+            if code == CloseCode.INTERNAL_ERROR.value and EXPIRED_TOKEN in reason:
                 logger.info("SAS token expired, refreshing...")
             else:
-                logger.warning(f"Connection closed with code {e.code}: {e.reason}")
-                raise e
+                logger.warning(f"Connection closed with code {code}: {reason}")
+                logger.warning("Retrying in 5 seconds...")
+                await asyncio.sleep(5)
         except Exception as e:
             logger.warning(f"Connection error: {e}")
             logger.warning("Retrying in 5 seconds...")

tests/test_relay_listener.py

@@ -138,12 +138,18 @@ async def test_main_handles_connection_closed_and_keyboard_interrupt(
 
     relay_listener_instance.listen.side_effect = [
         ConnectionClosedError(Close(CloseCode.INTERNAL_ERROR, "ExpiredToken"), None),
+        ConnectionClosedError(Close(CloseCode.INTERNAL_ERROR, "Something else"), None),
+        ConnectionClosedError(Close(CloseCode.BAD_GATEWAY, "Bad gateway"), None),
         KeyboardInterrupt(),
     ]
 
     await main()
 
-    assert relay_listener_instance.listen.call_count == 2
+    assert relay_listener_instance.listen.call_count == 4
     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("Connection closed with code 1011: Something else")
+    mock_logger.warning.assert_any_call("Retrying in 5 seconds...")
+    mock_logger.warning.assert_any_call("Connection closed with code 1014: Bad gateway")
+    mock_logger.warning.assert_any_call("Retrying in 5 seconds...")
     mock_logger.warning.assert_any_call("\nShutting down...")