Raise a token expiry specific error

Distinguish between TimeoutError and a token expired event by raising a custom error.

Author: steventux

src/relay_listener.py

@@ -45,6 +45,10 @@ class CredentialNotAvailableError(RuntimeError):
     pass
 
 
+class RelayTokenExpiredError(RuntimeError):
+    pass
+
+
 class RelayListener:
     """
     Socket Listener for Azure Relay.
@@ -86,15 +90,15 @@ async def listen(self):
                 async with self._connect(connection_url) as websocket:
                     logger.info("Connected - waiting for worklist actions...")
                     await self._listen_on_connection(websocket, refresh_at)
-            except asyncio.TimeoutError:
+            except RelayTokenExpiredError:
                 logger.info("Refreshing Azure Relay connection before expiry.")
                 continue
 
     async def _listen_on_connection(self, websocket, refresh_at: int):
         while True:
             timeout = refresh_at - time.time()
             if timeout <= 0:
-                raise asyncio.TimeoutError
+                raise RelayTokenExpiredError("Azure Relay token expired, refreshing.")
 
             try:
                 message = await asyncio.wait_for(websocket.recv(), timeout=timeout)

tests/test_relay_listener.py

@@ -8,7 +8,7 @@
 from websockets.frames import Close, CloseCode
 
 from models import WorklistItem
-from relay_listener import RelayListener, RelayURI, main, verify_credentials
+from relay_listener import RelayListener, RelayTokenExpiredError, RelayURI, main, verify_credentials
 
 
 class TestRelayListener:
@@ -40,7 +40,7 @@ async def test_listen_on_connection_echo(self, storage_instance):
         websocket = AsyncMock()
         websocket.recv.side_effect = [
             json.dumps({"accept": {"address": "wss://accept-url"}}),
-            asyncio.TimeoutError(),
+            RelayTokenExpiredError("Azure Relay token expired"),
         ]
 
         client_ws = AsyncMock()
@@ -51,7 +51,7 @@ async def test_listen_on_connection_echo(self, storage_instance):
         client_cm.__aexit__.return_value = None
 
         with patch("relay_listener.connect", return_value=client_cm) as mock_connect:
-            with pytest.raises(asyncio.TimeoutError):
+            with pytest.raises(RelayTokenExpiredError):
                 await subject._listen_on_connection(
                     websocket,
                     refresh_at=9999999999,
@@ -259,7 +259,7 @@ async def test_listen_refreshes_connection_after_timeout(
             patch.object(
                 subject,
                 "_listen_on_connection",
-                side_effect=[asyncio.TimeoutError(), KeyboardInterrupt()],
+                side_effect=[RelayTokenExpiredError(), KeyboardInterrupt()],
             ),
         ):
             with pytest.raises(KeyboardInterrupt):