🥅 Reauthenticate when websocket token expires

Author: shnizzedy

CHANGELOG.md

@@ -5,6 +5,12 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## 1.10.1
+
+### Fixed
+
+- Websocket now gets a new token when the one it's trying expires.
+
 ## 1.10.0
 
 ### Added

VERSION

@@ -1 +1 @@
-1.10.0
+1.10.1

package.json

@@ -1,6 +1,6 @@
 {
   "name": "hbnmigration",
-  "version": "1.10.0",
+  "version": "1.10.1",
   "private": true,
   "description": "HBN data migration monitoring infrastructure with Python and Node.js services",
   "workspaces": [

python_jobs/VERSION

@@ -1 +1 @@
-1.10.0
+1.10.1

python_jobs/src/hbnmigration/from_curious/alerts_to_redcap.py

@@ -388,7 +388,7 @@ async def websocket_listener(
 
 
 async def main_with_reconnect(
-    tokens: curious_variables.Tokens,
+    applet_name: str,
     uri: str,
     partial_redcap_landing: bool = False,
     max_attempts: Optional[int] = WS_MAX_RECONNECT_ATTEMPTS,
@@ -398,8 +398,8 @@ async def main_with_reconnect(
 
     Parameters
     ----------
-    tokens
-        Authentication tokens for WebSocket connection
+    applet_name
+        Name of applet to authenticate to
     uri
         WebSocket URI to connect to
     partial_redcap_landing
@@ -409,6 +409,8 @@ async def main_with_reconnect(
 
     """
     attempt = 0
+    tokens = curious_authenticate(applet_name)
+
     while max_attempts is None or attempt < max_attempts:
         try:
             if attempt > 0:
@@ -418,36 +420,43 @@ async def main_with_reconnect(
                     f" of {max_attempts}" if max_attempts else "",
                 )
             async with connect_to_websocket(tokens.access, uri) as websocket:
-                # Reset attempt counter on successful connection
                 if attempt > 0:
                     logger.info("Successfully reconnected to WebSocket")
                 attempt = 0
                 await websocket_listener(websocket, partial_redcap_landing)
                 logger.info("WebSocket listener completed normally")
                 break
+
         except ConnectionClosedError:
             attempt += 1
-            if max_attempts and attempt >= max_attempts:
+            if max_attempts is not None and attempt >= max_attempts:
                 logger.exception("Max reconnection attempts reached. Exiting.")
                 raise
             logger.warning(
                 "Connection lost. Reconnecting in %d seconds...", WS_RECONNECT_DELAY
             )
             await asyncio.sleep(WS_RECONNECT_DELAY)
+
         except InvalidStatus as e:
-            # Authentication or server errors
-            logger.exception(
-                "WebSocket connection failed with status %s", e.response.status_code
-            )
-            if e.response.status_code == requests.codes["unauthorized"]:
-                logger.exception(
-                    "Authentication failed. Token may be invalid or expired."
-                )
+            status = e.response.status_code
+            logger.exception("WebSocket connection failed with status %s", status)
+
+            if status == requests.codes["unauthorized"]:
+                # ── Re-authenticate ──
+                logger.warning("Token expired or invalid. Re-authenticating...")
+                try:
+                    tokens = curious_authenticate(applet_name)
+                    logger.info("Re-authentication successful")
+                except Exception:
+                    logger.exception("Re-authentication failed. Exiting.")
+                    raise
+
             attempt += 1
-            if max_attempts and attempt >= max_attempts:
+            if max_attempts is not None and attempt >= max_attempts:
                 logger.exception("Max reconnection attempts reached. Exiting.")
                 raise
             await asyncio.sleep(WS_RECONNECT_DELAY)
+
         except asyncio.CancelledError:
             logger.info("Operation cancelled")
             raise
@@ -485,10 +494,9 @@ async def main(
 
     """
     for applet_name in applet_names:
-        tokens = curious_authenticate(applet_name)
         endpoints = curious_variables.Endpoints(protocol="wss")
         await main_with_reconnect(
-            tokens=tokens,
+            applet_name=applet_name,
             uri=endpoints.alerts,
             partial_redcap_landing=partial_redcap_landing,
             max_attempts=max_attempts,

python_jobs/src/tests/conftest.py

@@ -15,7 +15,6 @@
 import requests
 from websockets.exceptions import InvalidStatus
 
-from hbnmigration._config_variables.curious_variables import Tokens
 from hbnmigration.from_curious.alerts_to_redcap import (
     parse_alert,
     synchronous_main,
@@ -698,16 +697,6 @@ def assert_pushed_data_contains_value(
     assert expected_value in pushed_data["value"].values
 
 
-def assert_reconnect_called_with_tokens_and_uri(
-    mock_reconnect: Mock, tokens: Tokens, uri_fragment: str
-) -> None:
-    """Assert reconnect was called with expected token and URI."""
-    assert mock_reconnect.called
-    call_kwargs = mock_reconnect.call_args[1]
-    assert call_kwargs["tokens"] == tokens
-    assert uri_fragment in call_kwargs["uri"]
-
-
 def assert_reconnect_called_with_max_attempts(
     mock_reconnect: Mock, expected_attempts: int
 ) -> None:
@@ -851,9 +840,13 @@ def setup_reconnect_mocks(
                     new_callable=AsyncMock,
                 )
             ),
+            "auth": stack.enter_context(
+                patch("hbnmigration.from_curious.alerts_to_redcap.curious_authenticate")
+            ),
         }
         mock_websocket = AsyncMock()
         mocks["ws"].return_value.__aenter__.return_value = mock_websocket
+        mocks["auth"].return_value = create_mock_tokens_ws()
         if listener_side_effect is not None:
             mocks["listener"].side_effect = listener_side_effect
         yield mocks

python_jobs/src/tests/test_alerts_curious_to_redcap.py

@@ -6,6 +6,7 @@
 
 import pandas as pd
 import pytest
+import requests
 from websockets.exceptions import ConnectionClosedError, InvalidStatus
 
 from hbnmigration.from_curious.alerts_to_redcap import (
@@ -345,6 +346,72 @@ async def test_websocket_listener_skips_non_answer_messages(mock_alerts_dependen
     assert not mock_alerts_dependencies["push"].called
 
 
+@pytest.mark.asyncio
+async def test_main_with_reconnect_reauths_on_401_then_succeeds(caplog):
+    """Test that main_with_reconnect re-authenticates on 401 and succeeds."""
+    with (
+        patch(
+            "hbnmigration.from_curious.alerts_to_redcap.connect_to_websocket"
+        ) as mock_ws,
+        patch(
+            "hbnmigration.from_curious.alerts_to_redcap.websocket_listener",
+            new_callable=AsyncMock,
+        ) as mock_listener,
+        patch(
+            "hbnmigration.from_curious.alerts_to_redcap.asyncio.sleep",
+            new_callable=AsyncMock,
+        ),
+        patch(
+            "hbnmigration.from_curious.alerts_to_redcap.curious_authenticate"
+        ) as mock_auth,
+    ):
+        mock_auth.return_value = create_mock_tokens_ws()
+        # First connect raises 401, second succeeds
+        mock_ctx = AsyncMock()
+        mock_ctx.__aenter__.return_value = AsyncMock()
+        mock_ws.side_effect = [
+            create_mock_invalid_status(requests.codes["unauthorized"]),
+            mock_ctx,
+        ]
+
+        await main_with_reconnect(
+            applet_name="test_applet", uri="wss://test.com/alerts", max_attempts=2
+        )
+        # Initial auth + re-auth on 401
+        assert mock_auth.call_count == 2
+        assert mock_listener.called
+        assert "Re-authentication successful" in caplog.text
+
+
+@pytest.mark.asyncio
+async def test_main_with_reconnect_reauth_failure_raises(caplog):
+    """Test that main_with_reconnect raises when re-authentication itself fails."""
+    with (
+        patch(
+            "hbnmigration.from_curious.alerts_to_redcap.connect_to_websocket"
+        ) as mock_ws,
+        patch(
+            "hbnmigration.from_curious.alerts_to_redcap.asyncio.sleep",
+            new_callable=AsyncMock,
+        ),
+        patch(
+            "hbnmigration.from_curious.alerts_to_redcap.curious_authenticate"
+        ) as mock_auth,
+    ):
+        # Initial auth succeeds, re-auth on 401 fails
+        mock_auth.side_effect = [
+            create_mock_tokens_ws(),
+            Exception("Auth server down"),
+        ]
+        mock_ws.side_effect = create_mock_invalid_status(requests.codes["unauthorized"])
+
+        with pytest.raises(Exception, match="Auth server down"):
+            await main_with_reconnect(
+                applet_name="test_applet", uri="wss://test.com/alerts", max_attempts=2
+            )
+        assert "Re-authentication failed" in caplog.text
+
+
 # ============================================================================
 # Tests - main_with_reconnect
 # ============================================================================
@@ -357,9 +424,8 @@ async def test_main_with_reconnect_successful_connection(
     """Test that main_with_reconnect handles successful connection."""
     setup_standard_alert_mocks(mock_alerts_dependencies, redcap_alerts_metadata)
     with setup_reconnect_mocks() as mocks:
-        tokens = create_mock_tokens_ws()
         await main_with_reconnect(
-            tokens=tokens, uri="wss://test.com/alerts", max_attempts=1
+            applet_name="test_applet", uri="wss://test.com/alerts", max_attempts=1
         )
         assert mocks["ws"].called
         assert mocks["listener"].called
@@ -369,9 +435,8 @@ async def test_main_with_reconnect_successful_connection(
 async def test_main_with_reconnect_handles_connection_error(caplog):
     """Test that main_with_reconnect reconnects on ConnectionClosedError."""
     with setup_reconnect_mocks([ConnectionClosedError(None, None), None]) as mocks:
-        tokens = create_mock_tokens_ws()
         await main_with_reconnect(
-            tokens=tokens, uri="wss://test.com/alerts", max_attempts=2
+            applet_name="test_applet", uri="wss://test.com/alerts", max_attempts=2
         )
         assert mocks["listener"].call_count == 2
         assert "Reconnecting in" in caplog.text
@@ -382,10 +447,9 @@ async def test_main_with_reconnect_handles_connection_error(caplog):
 async def test_main_with_reconnect_max_attempts_exceeded():
     """Test that main_with_reconnect stops after max attempts."""
     with setup_reconnect_mocks(ConnectionClosedError(None, None)) as mocks:
-        tokens = create_mock_tokens_ws()
         with pytest.raises(ConnectionClosedError):
             await main_with_reconnect(
-                tokens=tokens, uri="wss://test.com/alerts", max_attempts=1
+                applet_name="test_applet", uri="wss://test.com/alerts", max_attempts=1
             )
         assert mocks["listener"].call_count == 1
 
@@ -398,25 +462,34 @@ async def test_main_with_reconnect_handles_auth_error(caplog):
             "hbnmigration.from_curious.alerts_to_redcap.connect_to_websocket"
         ) as mock_ws,
         patch("hbnmigration.from_curious.alerts_to_redcap.asyncio.sleep") as mock_sleep,
+        patch(
+            "hbnmigration.from_curious.alerts_to_redcap.curious_authenticate"
+        ) as mock_auth,
     ):
-        mock_ws.side_effect = create_mock_invalid_status(401)
+        mock_ws.side_effect = create_mock_invalid_status(requests.codes["unauthorized"])
         mock_sleep.return_value = None
-        tokens = create_mock_tokens_ws()
+        # First call is the initial auth; second call is the re-auth on 401
+        mock_auth.side_effect = [
+            create_mock_tokens_ws(),
+            create_mock_tokens_ws(),
+        ]
         with pytest.raises(InvalidStatus):
             await main_with_reconnect(
-                tokens=tokens, uri="wss://test.com/alerts", max_attempts=1
+                applet_name="test_applet", uri="wss://test.com/alerts", max_attempts=1
             )
-        assert "Authentication failed" in caplog.text
+        assert (
+            "Token expired or invalid" in caplog.text
+            or "Authentication failed" in caplog.text
+        )
 
 
 @pytest.mark.asyncio
 async def test_main_with_reconnect_infinite_retries():
     """Test that main_with_reconnect runs indefinitely with None max_attempts."""
     side_effects = [ConnectionClosedError(None, None)] * 5 + [None]
     with setup_reconnect_mocks(side_effects) as mocks:
-        tokens = create_mock_tokens_ws()
         await main_with_reconnect(
-            tokens=tokens, uri="wss://test.com/alerts", max_attempts=None
+            applet_name="test_applet", uri="wss://test.com/alerts", max_attempts=None
         )
         assert mocks["listener"].call_count == 6
         assert mocks["sleep"].call_count == 5
@@ -448,22 +521,16 @@ async def test_main_processes_websocket_messages(
             await main(applet_names=["Healthy Brain Network Questionnaires"])
             assert mock_reconnect.called
             call_kwargs = mock_reconnect.call_args[1]
-            assert "tokens" in call_kwargs
+            assert "applet_name" in call_kwargs
             assert "wss://" in call_kwargs["uri"]
 
 
 @pytest.mark.asyncio
 async def test_main_passes_max_attempts():
     """Test that main() passes max_attempts parameter."""
-    with (
-        patch(
-            "hbnmigration.from_curious.alerts_to_redcap.curious_authenticate"
-        ) as mock_auth,
-        patch(
-            "hbnmigration.from_curious.alerts_to_redcap.main_with_reconnect"
-        ) as mock_reconnect,
-    ):
-        mock_auth.return_value = create_mock_tokens_ws()
+    with patch(
+        "hbnmigration.from_curious.alerts_to_redcap.main_with_reconnect"
+    ) as mock_reconnect:
         mock_reconnect.return_value = None
         await main(
             applet_names=["Healthy Brain Network Questionnaires"],