diff --git a/test/test_token_refresh.py b/test/test_token_refresh.py new file mode 100644 index 000000000..c40dbfbb8 --- /dev/null +++ b/test/test_token_refresh.py @@ -0,0 +1,175 @@ +"""Tests for the background OIDC token-refresh thread in `weaviate.connect.v4`. + +These run without a Weaviate instance and without an identity provider: the +`OAuth2Client` is real but its network-facing methods are replaced. +""" + +import time +from typing import Any, Dict, List, Optional + +import pytest +from authlib.integrations.base_client import OAuthError +from authlib.integrations.httpx_client import OAuth2Client +from httpx import HTTPError + +from weaviate.connect.v4 import ConnectionSync +from weaviate.warnings import _Warnings + +TOKEN_ENDPOINT = "https://localhost/token" + + +def test_oauth_error_is_not_an_http_error() -> None: + """Guards the reason the dedicated `except OAuthError` branch has to exist.""" + assert not issubclass(OAuthError, HTTPError) + + +class _FakeSession: + """Stands in for the authlib session returned by `_Auth.get_auth_session()`.""" + + def __init__(self, token: Dict[str, Any]) -> None: + self.__token = token + + def fetch_token(self) -> Dict[str, Any]: + return self.__token + + +class _FakeAuth: + """Minimal duck-typed stand-in for `_Auth`, which only needs these two methods.""" + + def __init__(self, token: Dict[str, Any]) -> None: + self.__token = token + self.calls = 0 + + def get_auth_session(self) -> "_FakeSession": + return _FakeSession(self.__token) + + def result(self, session: "_FakeSession") -> "_FakeSession": + self.calls += 1 + return session + + +def _make_client(token: Dict[str, Any]) -> OAuth2Client: + client = OAuth2Client(client_id="test-client", token=token) + client.metadata["token_endpoint"] = TOKEN_ENDPOINT + return client + + +def _make_connection(client: OAuth2Client) -> ConnectionSync: + """Build a connection with only the attributes the refresh helper touches. + + `_shutdown_background_event` is deliberately left unset: it is assigned by + `_create_background_token_refresh` itself before the thread reads it. + """ + connection = object.__new__(ConnectionSync) + connection._client = client + return connection + + +@pytest.fixture +def recorded_warnings(monkeypatch: pytest.MonkeyPatch) -> List[Exception]: + recorded: List[Exception] = [] + monkeypatch.setattr( + _Warnings, "token_refresh_failed", staticmethod(lambda exc: recorded.append(exc)) + ) + return recorded + + +def _run_refresh( + connection: ConnectionSync, auth: Optional[_FakeAuth], seconds: float = 3.0 +) -> None: + """Start the background refresh and let it complete at least one cycle.""" + connection._create_background_token_refresh(auth) # type: ignore[arg-type] + try: + time.sleep(seconds) + finally: + assert connection._shutdown_background_event is not None + connection._shutdown_background_event.set() + + +def _refresh_thread_is_alive() -> bool: + import threading + + return any( + thread.name == "TokenRefresh" and thread.is_alive() for thread in threading.enumerate() + ) + + +def test_oauth_error_does_not_kill_the_refresh_thread( + recorded_warnings: List[Exception], +) -> None: + """An `invalid_grant` from the IdP must be reported, not silently end all refreshes.""" + client = _make_client({"access_token": "old", "refresh_token": "stale", "expires_in": 1}) + + def _reject(**kwargs: Any) -> Dict[str, Any]: + raise OAuthError(error="invalid_grant", description="Refresh token is invalid") + + client.refresh_token = _reject # type: ignore[method-assign] + + auth = _FakeAuth({"access_token": "fresh", "refresh_token": "fresh-refresh", "expires_in": 60}) + connection = _make_connection(client) + + _run_refresh(connection, auth) + + assert recorded_warnings, "the OAuth2 rejection was swallowed instead of warned about" + assert isinstance(recorded_warnings[0], OAuthError) + # Re-authenticated from scratch rather than retrying the rejected refresh token. + assert auth.calls >= 1 + assert client.token["access_token"] == "fresh" + assert _refresh_thread_is_alive() + + +def test_oauth_error_without_credentials_still_keeps_the_thread_alive( + recorded_warnings: List[Exception], +) -> None: + """With no `_Auth` to fall back on, the loop must warn and keep running.""" + client = _make_client({"access_token": "old", "refresh_token": "stale", "expires_in": 1}) + + def _reject(**kwargs: Any) -> Dict[str, Any]: + raise OAuthError(error="invalid_grant", description="Refresh token is invalid") + + client.refresh_token = _reject # type: ignore[method-assign] + connection = _make_connection(client) + + _run_refresh(connection, None) + + assert recorded_warnings, "the OAuth2 rejection was swallowed instead of warned about" + assert all(isinstance(exc, OAuthError) for exc in recorded_warnings) + assert _refresh_thread_is_alive() + + +def test_http_error_still_retries(recorded_warnings: List[Exception]) -> None: + """The pre-existing `HTTPError` behaviour is unchanged.""" + client = _make_client({"access_token": "old", "refresh_token": "stale", "expires_in": 1}) + + def _fail(**kwargs: Any) -> Dict[str, Any]: + raise HTTPError("connection reset") + + client.refresh_token = _fail # type: ignore[method-assign] + + auth = _FakeAuth({"access_token": "fresh", "expires_in": 60}) + connection = _make_connection(client) + + _run_refresh(connection, auth) + + assert recorded_warnings + assert all(isinstance(exc, HTTPError) for exc in recorded_warnings) + # An HTTPError must not trigger the re-authentication path. + assert auth.calls == 0 + assert client.token["access_token"] == "old" + assert _refresh_thread_is_alive() + + +def test_successful_refresh_does_not_warn(recorded_warnings: List[Exception]) -> None: + """Sanity check that the happy path is untouched.""" + client = _make_client({"access_token": "old", "refresh_token": "good", "expires_in": 1}) + + def _succeed(**kwargs: Any) -> Dict[str, Any]: + return {"access_token": "new", "refresh_token": "good", "expires_in": 60} + + client.refresh_token = _succeed # type: ignore[method-assign] + connection = _make_connection(client) + + _run_refresh(connection, None) + + assert recorded_warnings == [] + assert client.token["access_token"] == "new" diff --git a/weaviate/connect/v4.py b/weaviate/connect/v4.py index d70bfc77b..40ce9a35e 100644 --- a/weaviate/connect/v4.py +++ b/weaviate/connect/v4.py @@ -23,6 +23,7 @@ ) import grpc +from authlib.integrations.base_client import OAuthError # type: ignore from authlib.integrations.httpx_client import ( # type: ignore AsyncOAuth2Client, OAuth2Client, @@ -595,6 +596,20 @@ def periodic_refresh_token(refresh_time: int, _auth: Optional[_Auth]) -> None: # retry again after one second, might be an unstable connection refresh_time = 1 _Warnings.token_refresh_failed(exc) + except OAuthError as exc: + # The identity provider rejected the refresh at the OAuth2 protocol level, + # e.g. the refresh token was rotated, reused or revoked. OAuthError is not an + # HTTPError, so without this branch it would escape the thread and silently + # stop all further refreshes. Retrying the same refresh token cannot succeed, + # so re-authenticate from scratch when credentials are available. + refresh_time = 1 + _Warnings.token_refresh_failed(exc) + if _auth is not None: + try: + refresh_session() + refresh_time = update_refresh_time() + except (HTTPError, OAuthError) as reauth_exc: + _Warnings.token_refresh_failed(reauth_exc) demon = Thread( target=periodic_refresh_token,